forked from hofmannol/AlgoDatSoSe25
60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import random
|
|
import pygame
|
|
from utils.game import Game
|
|
from avl_tree import AVLTree
|
|
|
|
WHITE = (255, 255, 255)
|
|
BLUE = (0, 0, 255)
|
|
BLACK = (0, 0, 0)
|
|
WIDTH = 800
|
|
HEIGHT = 400
|
|
MARGIN = 20
|
|
|
|
class AVLTreeGame(Game):
|
|
|
|
def __init__(self):
|
|
super().__init__("AVLTree Game", fps=10, size=(WIDTH, HEIGHT))
|
|
random.seed()
|
|
self.z = list(range(1, 501))
|
|
random.shuffle(self.z)
|
|
self.finished = False
|
|
self.tree = AVLTree()
|
|
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)
|
|
self.generator = None
|
|
|
|
def update_game(self):
|
|
if not self.finished:
|
|
if self.generator is None:
|
|
self.generator = self.tree.insert_stepwise(self.z.pop())
|
|
try:
|
|
next(self.generator)
|
|
except StopIteration:
|
|
self.generator = None
|
|
if self.generator is None and len(self.z) == 0:
|
|
self.finished = True
|
|
self.height = self.tree.get_height(self.tree.root)
|
|
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 = AVLTreeGame()
|
|
tree_game.run()
|
|
|