Implementation of Red & Black Knights - Numberphile by Claude

red_black_knights.py

"""
Red & Black Knights on a square spiral.

Construction (OEIS A392177 / A392178 / A395357, Numberphile 2026-05-12):

    Number the cells of an infinite square spiral starting at 0 in the center,
    moving right first and then counterclockwise (OEIS A274641 layout).
    Black and Red take turns, Black first.
        - On Black's turn, place a knight on the smallest unoccupied cell
          that is NOT attacked by an existing Red knight.
        - On Red's turn, same rule with colors swapped.
    Same-color knights are allowed to attack each other; only the OPPOSITE
    color matters.

The sequences:
    A392177 = cells occupied by Black knights, in cell-index order
    A392178 = cells occupied by Red knights, in cell-index order
    A395357 = the final spiral, read in cell order, with +n for Black, -n for
              Red, 0 for unoccupied (cell 0 itself is encoded as 0 by
              convention even though Black sits there).
"""

import math
from collections import namedtuple

# Knight move offsets (the 8 L-moves).
KNIGHT_OFFSETS = [
    (1, 2), (1, -2), (-1, 2), (-1, -2),
    (2, 1), (2, -1), (-2, 1), (-2, -1),
]


# --------------------------------------------------------------------------
# Spiral coordinates
# --------------------------------------------------------------------------
def cell_to_xy(n):
    """Map a non-negative cell index n to its (x, y) coordinate on the
    square spiral. Cell 0 is at the origin; cell 1 is at (1, 0); the spiral
    then proceeds counterclockwise.

    Shell k (k >= 1) occupies cells (2k-1)^2 .. (2k+1)^2 - 1, total 8k cells,
    starting at (k, -(k-1)) and walking up, left, down, right around the
    bounding square of side 2k+1.
    """
    if n == 0:
        return (0, 0)
    # k is the Chebyshev radius (the "shell"): k = max(|x|, |y|).
    k = math.isqrt(n - 1) // 2 + 1   # initial guess
    # Tighten in case isqrt was off.
    while (2 * k - 1) ** 2 > n:
        k -= 1
    while (2 * k + 1) ** 2 - 1 < n:
        k += 1
    m = n - (2 * k - 1) ** 2          # 0 <= m < 8k
    side, pos = divmod(m, 2 * k)      # which of the 4 sides, and how far along
    if side == 0:   # right side, walking up
        return (k, -k + 1 + pos)
    if side == 1:   # top side, walking left
        return (k - 1 - pos, k)
    if side == 2:   # left side, walking down
        return (-k, k - 1 - pos)
    return (-k + 1 + pos, -k)         # bottom side, walking right


# --------------------------------------------------------------------------
# Simulation
# --------------------------------------------------------------------------
SimResult = namedtuple("SimResult", ["color_of_cell", "xy_of_cell", "black_cells", "red_cells"])


def simulate(num_moves):
    """Run the game for `num_moves` half-moves (so num_moves // 2 Black moves
    and num_moves // 2 Red moves, give or take one). Returns a SimResult.

    Implementation notes:
      - We keep per-color "attack sets" as sets of (x, y) tuples. Adding a
        knight is O(1) (8 offsets); checking attack is O(1) hash lookup.
      - We keep a per-color pointer into the cell-index sequence. The pointer
        only ever advances forward: once a cell is occupied or attacked by
        the opponent, it never becomes available to that player again.
    """
    color_of_cell = {}     # cell index -> 'B' or 'R'
    xy_of_cell = {}        # cell index -> (x, y)
    black_attacks = set()  # (x, y) attacked by some Black knight
    red_attacks = set()    # (x, y) attacked by some Red knight
    black_cells = []
    red_cells = []

    # Smallest cell index this color might still legally take.
    pointer = {'B': 0, 'R': 0}

    for turn in range(num_moves):
        color = 'B' if (turn % 2 == 0) else 'R'
        enemy_attacks = red_attacks if color == 'B' else black_attacks

        c = pointer[color]
        while True:
            if c not in color_of_cell:
                xy = cell_to_xy(c)
                if xy not in enemy_attacks:
                    break
            c += 1

        # Place the knight.
        color_of_cell[c] = color
        xy_of_cell[c] = xy
        (black_cells if color == 'B' else red_cells).append(c)
        my_attacks = black_attacks if color == 'B' else red_attacks
        for dx, dy in KNIGHT_OFFSETS:
            my_attacks.add((xy[0] + dx, xy[1] + dy))

        # This color's pointer moves past c. We don't try to advance further;
        # next turn we'll skip over whatever's now blocked.
        pointer[color] = c + 1

    return SimResult(color_of_cell, xy_of_cell, black_cells, red_cells)


# --------------------------------------------------------------------------
# OEIS spot-check
# --------------------------------------------------------------------------
# A392177, the first 75 black-knight cells (from the OEIS listing).
A392177_HEAD = [
    0, 2, 5, 9, 11, 15, 20, 21, 30, 31, 36, 40, 42, 47, 48, 50, 56, 61, 65,
    67, 69, 70, 71, 75, 76, 81, 83, 85, 87, 89, 93, 99, 109, 110, 111, 112,
    116, 117, 126, 132, 133, 138, 144, 148, 150, 152, 154, 156, 161, 162,
    176, 180, 182, 187, 193, 197, 199, 201, 203, 205, 207, 208, 209, 211,
    213, 214, 219, 229, 231, 233, 235, 237, 238, 239, 243,
]

# A395357, the final spiral in cell order (heading values from the OEIS).
A395357_HEAD = [
    0, -1, 2, -3, -4, 5, -6, 0, 0, 9, -10, 11, -12, 0, 0, 15, 0, 0, 0, 0,
    20, 21, 0, 0, -24, -25, 0, 0, 0, 0, 30, 31, 0, 0, -34, -35, 36, -37,
    0, 0, 40, -41, 42, 0, -44, 0, 0, 47, 48, -49, 50, 0, 0, 0, 0, -55, 56,
    -57, -58, 0, 0, 61, 0, -63, -64, 65, -66, 67, -68, 69, 70, 71, -72, 0,
    0, 75, 76, 0, -78, 0, 0, 81, -82, 83, -84, 85, -86, 87, -88,
]


def verify_against_oeis():
    # Need enough turns to reach cell index 243 (the largest in A392177_HEAD).
    # Each turn lays one knight, so 2 * (largest_cell + 1) is plenty.
    sim = simulate(2 * 250)
    ours = sorted(sim.black_cells)[: len(A392177_HEAD)]
    assert ours == A392177_HEAD, f"A392177 mismatch:\n  ours: {ours}\n  oeis: {A392177_HEAD}"

    # A395357 is the cell-indexed view: +c for Black, -c for Red, 0 unoccupied
    # (with cell 0 reported as 0 by convention).
    spiral = []
    for c in range(len(A395357_HEAD)):
        if c == 0:
            spiral.append(0)
        elif c in sim.color_of_cell:
            spiral.append(c if sim.color_of_cell[c] == 'B' else -c)
        else:
            spiral.append(0)
    assert spiral == A395357_HEAD, f"A395357 mismatch:\n  ours: {spiral}\n  oeis: {A395357_HEAD}"
    print("OEIS spot-check passed: A392177 head + A395357 head both match.")


# --------------------------------------------------------------------------
# Visualization
# --------------------------------------------------------------------------
def plot_spiral(sim, out_path, title=None, figsize=10, dpi=200):
    """Rasterize the colored cells to a numpy array and imshow it. This is
    far faster and crisper than scatter for tens or hundreds of thousands of
    cells: every occupied cell becomes exactly one pixel of the right color,
    and unoccupied cells stay white."""
    import numpy as np
    import matplotlib.pyplot as plt

    radius = max(max(abs(x), abs(y)) for x, y in sim.xy_of_cell.values())
    side = 2 * radius + 1
    # white background; we paint cells black or red as we go.
    img = np.full((side, side, 3), 255, dtype=np.uint8)
    for c, (x, y) in sim.xy_of_cell.items():
        # image[row, col]; image row 0 is at the top, so flip y.
        row = radius - y
        col = radius + x
        if sim.color_of_cell[c] == 'B':
            img[row, col] = (17, 17, 17)
        else:
            img[row, col] = (200, 30, 30)

    fig, ax = plt.subplots(figsize=(figsize, figsize), facecolor="white")
    ax.imshow(img, interpolation="nearest")
    ax.axhline(radius, color="#888", linewidth=0.4)
    ax.axvline(radius, color="#888", linewidth=0.4)
    ax.set_xticks([]); ax.set_yticks([])
    for spine in ax.spines.values():
        spine.set_visible(False)
    if title:
        ax.set_title(title, fontsize=11)
    fig.tight_layout()
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    print(f"wrote {out_path}   "
          f"(radius {radius}, Black {len(sim.black_cells)}, Red {len(sim.red_cells)})")


if __name__ == "__main__":
    verify_against_oeis()

    for N in (20_000, 200_000):
        sim = simulate(N)
        plot_spiral(
            sim,
            f"/home/claude/red_black_knights_{N}.png",
            title=f"Red & Black Knights on the square spiral, first {N:,} moves",
        )

matplotlib==3.10.9
numpy==2.4.6

Source video: https://www.youtube.com/watch?v=UiX4CFIiegM