58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
import random
|
|
import pygame
|
|
from utils.game import Game
|
|
from utils.algo_context import AlgoContext
|
|
from bin_tree import BinaryTree
|
|
|
|
WHITE = (255, 255, 255)
|
|
BLUE = (0, 0, 255)
|
|
BLACK = (0, 0, 0)
|
|
WIDTH, HEIGHT, MARGIN = 800, 400, 20
|
|
|
|
|
|
class BinTreeGame(Game):
|
|
|
|
def __init__(self):
|
|
super().__init__("BinTree Game", fps=10, size=(WIDTH, HEIGHT))
|
|
random.seed()
|
|
self.z = list(range(1, 101))
|
|
random.shuffle(self.z)
|
|
self.finished = False
|
|
self.ctx = AlgoContext()
|
|
self.tree = BinaryTree(self.ctx)
|
|
self.tree.get_height = lambda node: (
|
|
0 if node is None
|
|
else 1 + max(self.tree.get_height(node.left), self.tree.get_height(node.right))
|
|
)
|
|
self.height = self.tree.get_height(self.tree.root)
|
|
|
|
def update_game(self):
|
|
if not self.finished:
|
|
self.tree.insert(self.z.pop())
|
|
self.height = self.tree.get_height(self.tree.root)
|
|
if len(self.z) == 0:
|
|
self.finished = True
|
|
return True
|
|
|
|
def draw_game(self):
|
|
self.screen.fill(WHITE)
|
|
if self.height > 0:
|
|
self.draw_tree(self.tree.root, WIDTH // 2, MARGIN, WIDTH // 4 - MARGIN)
|
|
super().draw_game()
|
|
|
|
def draw_tree(self, node, x, y, x_offset):
|
|
y_offset = (HEIGHT - 2 * MARGIN) / self.height
|
|
if node is not None:
|
|
pygame.draw.circle(self.screen, BLUE, (x, y), 2)
|
|
if node.left is not None:
|
|
pygame.draw.line(self.screen, BLACK, (x, y), (x - x_offset, y + y_offset))
|
|
self.draw_tree(node.left, x - x_offset, y + y_offset, x_offset // 2)
|
|
if node.right is not None:
|
|
pygame.draw.line(self.screen, BLACK, (x, y), (x + x_offset, y + y_offset))
|
|
self.draw_tree(node.right, x + x_offset, y + y_offset, x_offset // 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tree_game = BinTreeGame()
|
|
tree_game.run()
|