Pythonの関数とクラスの復習

HanoiRod: TypeAlias = list[int]
HanoiState: TypeAlias = list[HanoiRod]


def visualize_hanoi(
    states: Sequence[Sequence[Sequence[int]]],
    *,
    interval: int = 600,
    repeat: bool = False,
    as_html: bool = True,
) -> animation.FuncAnimation | HTML:
    """ハノイの塔の盤面列をアニメーションとして表示する．

    Args:
        states: 盤面の列です．各盤面は 3 本の棒の状態を表し，
            各棒は下から上の順で円盤番号を並べたリストにします．
            例: ``[[3, 2], [1], []]``
        interval: コマ送り間隔をミリ秒で指定します．
        repeat: 最後まで再生したあとに繰り返すかどうかを指定します．
        as_html: ``True`` のとき，Notebook でそのまま表示しやすい
            ``IPython.display.HTML`` を返します．

    Returns:
        ``as_html=True`` なら ``HTML``，それ以外なら
        ``matplotlib.animation.FuncAnimation`` を返します．
    """
    if not states:
        raise ValueError("states must not be empty")

    normalized_states: list[HanoiState] = []
    for state in states:
        if len(state) != 3:
            raise ValueError("each state must have exactly three rods")
        normalized_state = [list(rod) for rod in state]
        normalized_states.append(normalized_state)

    disk_values = [disk for state in normalized_states for rod in state for disk in rod]
    if not disk_values:
        raise ValueError("states must contain at least one disk")
    if any(disk <= 0 for disk in disk_values):
        raise ValueError("disk numbers must be positive integers")

    max_disk = max(disk_values)
    max_height = max(len(rod) for state in normalized_states for rod in state)

    fig, ax = plt.subplots(figsize=(8, 4.5))
    fig.patch.set_facecolor("white")

    rod_x_positions = [0, 1, 2]
    rod_width = 0.08
    base_y = 0.0
    disk_height = 0.6
    disk_base_width = 0.28
    disk_width_step = 0.24
    colors = plt.cm.viridis(
        [i / max(max_disk - 1, 1) for i in range(max_disk)]
    )

    def setup_axes() -> None:
        ax.clear()
        ax.set_xlim(-0.6, 2.6)
        ax.set_ylim(-0.2, max_height * disk_height + 1.5)
        ax.set_xticks(rod_x_positions)
        ax.set_xticklabels(["左", "中央", "右"])
        ax.set_yticks([])
        ax.set_title("ハノイの塔")
        for spine in ax.spines.values():
            spine.set_visible(False)

        ax.add_patch(
            Rectangle(
                (-0.5, base_y - 0.15),
                3.0,
                0.15,
                color="#6b4f2a",
            )
        )

        for x in rod_x_positions:
            ax.add_patch(
                Rectangle(
                    (x - rod_width / 2, base_y),
                    rod_width,
                    max_height * disk_height + 0.3,
                    color="#8b6f47",
                )
            )

    def draw_state(state_index: int) -> list[Any]:
        setup_axes()
        state = normalized_states[state_index]
        artists: list[Any] = []

        for rod_index, rod in enumerate(state):
            x_center = rod_x_positions[rod_index]
            for level, disk in enumerate(rod):
                disk_width = disk_base_width + disk * disk_width_step
                x_left = x_center - disk_width / 2
                y_bottom = base_y + level * disk_height
                color = colors[disk - 1]
                rect = Rectangle(
                    (x_left, y_bottom),
                    disk_width,
                    disk_height * 0.82,
                    facecolor=color,
                    edgecolor="black",
                    linewidth=1.0,
                    zorder=3,
                )
                ax.add_patch(rect)
                artists.append(rect)
                label = ax.text(
                    x_center,
                    y_bottom + disk_height * 0.45,
                    str(disk),
                    ha="center",
                    va="center",
                    fontsize=10,
                    color="white",
                    zorder=4,
                )
                artists.append(label)

        step_label = ax.text(
            2.45,
            max_height * disk_height + 0.8,
            f"step = {state_index}",
            ha="right",
            va="center",
            fontsize=11,
        )
        artists.append(step_label)
        return artists

    ani = animation.FuncAnimation(
        fig,
        draw_state,
        frames=len(normalized_states),
        interval=interval,
        blit=False,
        repeat=repeat,
    )
    plt.close(fig)

    if as_html:
        #with open("ani.html", "w") as f:
        #    f.write(ani.to_jshtml())
        return HTML(ani.to_jshtml())
    return ani


example_states = [
    [[3, 2, 1], [], []],
    [[3, 2], [], [1]],
    [[3], [2], [1]],
    [[3], [2, 1], []],
    [[], [2, 1], [3]],
]

#visualize_hanoi(example_states)

LifeGameBoard = list[list[int]] # ライフゲームの盤面を表す型

def visualize_lifegame(
    states: Sequence[LifeGameBoard],
    *, # キーワード引数
    interval: int = 250,
    repeat: bool = False,
    as_html: bool = True,
) -> animation.FuncAnimation | HTML:
    """ライフゲームの盤面列をアニメーション表示する．

    Args:
        states: 各時刻の盤面を表す 2 次元配列の列です．
        interval: コマ送り間隔をミリ秒で指定します．
        repeat: 最後まで再生したあとに繰り返すかどうかを指定します．
        as_html: ``True`` のとき，Notebook で表示しやすい ``HTML`` を返します．

    Returns:
        ``as_html=True`` なら ``HTML``，それ以外なら
        ``matplotlib.animation.FuncAnimation`` を返します．
    """
    if not states:
        raise ValueError("states must not be empty")

    normalized_states: list[LifeGameBoard] = []
    row_count = len(states[0])
    col_count = len(states[0][0]) if row_count > 0 else 0
    if row_count == 0 or col_count == 0:
        raise ValueError("each state must be a non-empty 2D grid")

    for state in states:
        if len(state) != row_count:
            raise ValueError("all states must have the same number of rows")
        normalized_state = [list(row) for row in state]
        if any(len(row) != col_count for row in normalized_state):
            raise ValueError("all rows must have the same length")
        normalized_states.append(normalized_state)

    fig, ax = plt.subplots(figsize=(6, 6))
    image = ax.imshow(normalized_states[0], cmap="Greys", vmin=0, vmax=1)
    ax.set_title("ライフゲーム")
    ax.set_xticks([])
    ax.set_yticks([])

    step_text = ax.text(
        0.98,
        1.02,
        "step = 0",
        transform=ax.transAxes,
        ha="right",
        va="bottom",
        fontsize=11,
    )

    def update(frame_index: int) -> list[Any]:
        image.set_data(normalized_states[frame_index])
        step_text.set_text(f"step = {frame_index}")
        return [image, step_text]

    ani = animation.FuncAnimation(
        fig,
        update,
        frames=len(normalized_states),
        interval=interval,
        blit=False,
        repeat=repeat,
    )
    plt.close(fig)

    if as_html:
        #with open("lifegame.html", "w") as f:
        #    f.write(ani.to_jshtml())
        return HTML(ani.to_jshtml())
    return ani