some new stuff.
idk its all pretty fun! some C++ too!
This commit is contained in:
199
simulations/astar/main.py
Normal file
199
simulations/astar/main.py
Normal file
@@ -0,0 +1,199 @@
|
||||
import pygame
|
||||
import random
|
||||
from queue import PriorityQueue
|
||||
import time
|
||||
|
||||
WIDTH = 600
|
||||
ROWS = 20
|
||||
CELL_SIZE = WIDTH // ROWS
|
||||
FPS = 60
|
||||
DELAY = 0
|
||||
|
||||
WHITE = (255, 255, 255)
|
||||
BLACK = (0, 0, 0)
|
||||
GREEN = (0, 255, 0)
|
||||
RED = (255, 0, 0)
|
||||
BLUE = (0, 0, 255)
|
||||
GRAY = (200, 200, 200)
|
||||
YELLOW = (255, 255, 0)
|
||||
|
||||
class Cell:
|
||||
def __init__(self, row, col):
|
||||
self.row = row
|
||||
self.col = col
|
||||
self.x = col * CELL_SIZE
|
||||
self.y = row * CELL_SIZE
|
||||
self.walls = [True, True, True, True] # top, right, bottom, left
|
||||
self.visited = False
|
||||
|
||||
def draw(self, win):
|
||||
if self.visited:
|
||||
pygame.draw.rect(win, WHITE, (self.x, self.y, CELL_SIZE, CELL_SIZE))
|
||||
if self.walls[0]: pygame.draw.line(win, BLACK, (self.x, self.y), (self.x+CELL_SIZE, self.y), 2)
|
||||
if self.walls[1]: pygame.draw.line(win, BLACK, (self.x+CELL_SIZE, self.y), (self.x+CELL_SIZE, self.y+CELL_SIZE), 2)
|
||||
if self.walls[2]: pygame.draw.line(win, BLACK, (self.x+CELL_SIZE, self.y+CELL_SIZE), (self.x, self.y+CELL_SIZE), 2)
|
||||
if self.walls[3]: pygame.draw.line(win, BLACK, (self.x, self.y+CELL_SIZE), (self.x, self.y), 2)
|
||||
|
||||
def highlight(self, win, color):
|
||||
pygame.draw.rect(win, color, (self.x, self.y, CELL_SIZE, CELL_SIZE))
|
||||
|
||||
|
||||
def generate_maze(grid, current):
|
||||
stack = []
|
||||
current.visited = True
|
||||
|
||||
while True:
|
||||
neighbors = []
|
||||
r, c = current.row, current.col
|
||||
if r > 0 and not grid[r-1][c].visited: neighbors.append(grid[r-1][c])
|
||||
if r < ROWS-1 and not grid[r+1][c].visited: neighbors.append(grid[r+1][c])
|
||||
if c > 0 and not grid[r][c-1].visited: neighbors.append(grid[r][c-1])
|
||||
if c < ROWS-1 and not grid[r][c+1].visited: neighbors.append(grid[r][c+1])
|
||||
|
||||
if neighbors:
|
||||
next_cell = random.choice(neighbors)
|
||||
stack.append(current)
|
||||
remove_walls(current, next_cell)
|
||||
next_cell.visited = True
|
||||
current = next_cell
|
||||
elif stack:
|
||||
current = stack.pop()
|
||||
else:
|
||||
break
|
||||
|
||||
def remove_walls(a, b):
|
||||
dx = a.col - b.col
|
||||
dy = a.row - b.row
|
||||
if dx == 1:
|
||||
a.walls[3] = False
|
||||
b.walls[1] = False
|
||||
elif dx == -1:
|
||||
a.walls[1] = False
|
||||
b.walls[3] = False
|
||||
if dy == 1:
|
||||
a.walls[0] = False
|
||||
b.walls[2] = False
|
||||
elif dy == -1:
|
||||
a.walls[2] = False
|
||||
b.walls[0] = False
|
||||
|
||||
def heuristic(a, b):
|
||||
return abs(a.row - b.row) + abs(a.col - b.col)
|
||||
|
||||
def astar(grid, start, end, win):
|
||||
count = 0
|
||||
open_set = PriorityQueue()
|
||||
open_set.put((0, count, start))
|
||||
came_from = {}
|
||||
g_score = {cell: float("inf") for row in grid for cell in row}
|
||||
g_score[start] = 0
|
||||
f_score = {cell: float("inf") for row in grid for cell in row}
|
||||
f_score[start] = heuristic(start, end)
|
||||
open_set_hash = {start}
|
||||
|
||||
while not open_set.empty():
|
||||
for event in pygame.event.get():
|
||||
if event.type == pygame.QUIT:
|
||||
pygame.quit()
|
||||
return []
|
||||
|
||||
current = open_set.get()[2]
|
||||
open_set_hash.remove(current)
|
||||
|
||||
if current == end:
|
||||
path = []
|
||||
while current in came_from:
|
||||
current = came_from[current]
|
||||
path.append(current)
|
||||
return path
|
||||
|
||||
for neighbor in get_neighbors(grid, current):
|
||||
temp_g = g_score[current] + 1
|
||||
if temp_g < g_score[neighbor]:
|
||||
came_from[neighbor] = current
|
||||
g_score[neighbor] = temp_g
|
||||
f_score[neighbor] = temp_g + heuristic(neighbor, end)
|
||||
if neighbor not in open_set_hash:
|
||||
count += 1
|
||||
open_set.put((f_score[neighbor], count, neighbor))
|
||||
open_set_hash.add(neighbor)
|
||||
|
||||
# visualize
|
||||
draw_grid(win, grid)
|
||||
for cell in open_set_hash:
|
||||
cell.highlight(win, GREEN)
|
||||
pygame.display.update()
|
||||
time.sleep(DELAY)
|
||||
|
||||
return []
|
||||
|
||||
def get_neighbors(grid, cell):
|
||||
neighbors = []
|
||||
r, c = cell.row, cell.col
|
||||
if not cell.walls[0] and r > 0: neighbors.append(grid[r-1][c])
|
||||
if not cell.walls[1] and c < ROWS-1: neighbors.append(grid[r][c+1])
|
||||
if not cell.walls[2] and r < ROWS-1: neighbors.append(grid[r+1][c])
|
||||
if not cell.walls[3] and c > 0: neighbors.append(grid[r][c-1])
|
||||
return neighbors
|
||||
|
||||
# --- DRAWING ---
|
||||
def draw_grid(win, grid):
|
||||
win.fill(WHITE)
|
||||
for row in grid:
|
||||
for cell in row:
|
||||
cell.draw(win)
|
||||
|
||||
# --- MAIN ---
|
||||
def main():
|
||||
global ROWS, CELL_SIZE
|
||||
pygame.init()
|
||||
win = pygame.display.set_mode((WIDTH, WIDTH))
|
||||
pygame.display.set_caption("A*")
|
||||
|
||||
clock = pygame.time.Clock()
|
||||
|
||||
def create_maze(rows):
|
||||
global CELL_SIZE
|
||||
CELL_SIZE = WIDTH // rows
|
||||
grid = [[Cell(r, c) for c in range(rows)] for r in range(rows)]
|
||||
generate_maze(grid, grid[0][0])
|
||||
return grid
|
||||
|
||||
grid = create_maze(ROWS)
|
||||
start = grid[0][0]
|
||||
end = grid[ROWS-1][ROWS-1]
|
||||
path = []
|
||||
|
||||
running = True
|
||||
while running:
|
||||
clock.tick(FPS)
|
||||
draw_grid(win, grid)
|
||||
start.highlight(win, BLUE)
|
||||
end.highlight(win, RED)
|
||||
for cell in path:
|
||||
cell.highlight(win, YELLOW)
|
||||
pygame.display.update()
|
||||
|
||||
for event in pygame.event.get():
|
||||
if event.type == pygame.QUIT:
|
||||
running = False
|
||||
if event.type == pygame.KEYDOWN:
|
||||
if event.key == pygame.K_SPACE:
|
||||
path = astar(grid, start, end, win)
|
||||
if event.key == pygame.K_UP:
|
||||
ROWS += 1
|
||||
grid = create_maze(ROWS)
|
||||
start = grid[0][0]
|
||||
end = grid[ROWS-1][ROWS-1]
|
||||
path = []
|
||||
if event.key == pygame.K_DOWN and ROWS > 5:
|
||||
ROWS -= 1
|
||||
grid = create_maze(ROWS)
|
||||
start = grid[0][0]
|
||||
end = grid[ROWS-1][ROWS-1]
|
||||
path = []
|
||||
|
||||
pygame.quit()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user