85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
import matplotlib.pyplot as plt
|
|
import matplotlib.animation as animation
|
|
from time import sleep
|
|
import copy
|
|
import argparse
|
|
|
|
|
|
def animate_hanoi(n, frame_delay=750, src=0, dest=1, aux=2):
|
|
colors = []
|
|
for i in range(n):
|
|
# Calculate the hue value
|
|
hue = i / n
|
|
# Convert HSL to RGB
|
|
rgb = plt.cm.hsv(hue)[:3] # Use HSV colormap and take RGB values
|
|
colors.append(rgb)
|
|
|
|
fig, ax = plt.subplots()
|
|
fig.suptitle(f"Towers of Hanoi simulation for n={n}", fontsize=21)
|
|
ax.set_xlim(-1, 3)
|
|
ax.set_ylim(0, n + 1)
|
|
ax.set_xticks([0, 1, 2])
|
|
ax.set_xticklabels(['A', 'B', 'C'])
|
|
ax.set_yticks([])
|
|
|
|
frames = []
|
|
|
|
def move_and_frame(src, dest):
|
|
disk = towers[src].pop()
|
|
towers[dest].append(disk)
|
|
|
|
# Draw the disks
|
|
frame = []
|
|
for i, tower in enumerate(towers):
|
|
for j, disk in enumerate(tower):
|
|
frame += ax.bar(i, 1, bottom=j,
|
|
width=0.5, color=colors[disk-1]).patches
|
|
frame.append(ax.text(i, j + 0.5, str(disk), ha='center',
|
|
va='center', color='white', fontsize=16))
|
|
|
|
frames.append(frame)
|
|
|
|
def solve_game(n, src, dest, aux):
|
|
if n == 1:
|
|
move_and_frame(src, dest)
|
|
else:
|
|
solve_game(n - 1, src, aux, dest)
|
|
move_and_frame(src, dest)
|
|
solve_game(n - 1, aux, dest, src)
|
|
|
|
# Create the towers
|
|
towers = [[], [], []]
|
|
for i in range(n):
|
|
towers[src].append(n - i)
|
|
|
|
# Frame the initial state
|
|
move_and_frame(0, 0)
|
|
|
|
# Solve it
|
|
solve_game(n, src, dest, aux)
|
|
|
|
ani = animation.ArtistAnimation(fig, frames, interval=frame_delay, blit=True,
|
|
repeat=False)
|
|
|
|
return ani
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="A simple program to animate the optimal solution for a Towers of Hanoi game with n disks")
|
|
parser.add_argument("--disk-number", "-n", type=int,
|
|
help="Number of disks to simulate. Asks by default", default=0)
|
|
parser.add_argument("--interval", "-i", type=int,
|
|
help="Interval (in ms) between frames", default=750)
|
|
parser.add_argument("--output", "-o", type=str,
|
|
help="File to write the animation to (if specified)", default=None)
|
|
args = parser.parse_args()
|
|
if args.disk_number == 0:
|
|
args.disk_number = int(input("❔ Number of disks in the game: "))
|
|
ani = animate_hanoi(args.disk_number, args.interval)
|
|
if args.output:
|
|
ani.save(args.output)
|
|
print("Animation saved succesfully")
|
|
else:
|
|
plt.show()
|