checkerboard.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

def color_grid(a):
    n = a * a  # grid is n x n
    grid = np.full((n, n), -1, dtype=int)  # -1 means uncolored

    # Assign diagonal block colors
    for k in range(a):
        rs = k*a
        cs = k*a
        grid[rs:rs+a, cs:cs+a] = k

    # Handle off-diagonal blocks
    for br in range(a):
        for bc in range(a):
            if br == bc:
                continue  # diagonal, already filled
            rs, cs = br*a, bc*a
            block = np.zeros((a, a), dtype=int)
            # Case 1: Even a, use finer checkerboard in block
            if a % 2 == 0:
                for i in range(a):
                    for j in range(a):
                        if (i + j) % 2 == 0:
                            # 'row cell': color of the block's row-diagonal
                            grid[rs+i, cs+j] = br
                        else:
                            # 'col cell': color of the block's col-diagonal
                            grid[rs+i, cs+j] = bc
            # Case 2: Odd a, use checkerboard at block level
            else:
                # Determine "row" or "col" block
                if ((br + bc) % 2 == 0) == (br < bc):
                    # row cell: row block's diag color
                    grid[rs:rs+a, cs:cs+a] = br
                else:
                    # col cell: col block's diag color
                    grid[rs:rs+a, cs:cs+a] = bc

    return grid

def plot_color_grid(grid):
    a = int(np.sqrt(grid.shape[0]))
    fig, ax = plt.subplots(figsize=(8, 8))
    # Use tab20, tab10 or a custom colormap for more colors
    ncolors = np.max(grid) + 1
    colors = plt.cm.get_cmap('tab20', ncolors)
    norm = mcolors.BoundaryNorm(np.arange(-0.5, ncolors+0.5), ncolors)
    ax.imshow(grid, cmap=colors, norm=norm)
    # Draw grid lines for blocks
    for i in range(0, grid.shape[0]+1, a):
        ax.axhline(i-0.5, color='black', lw=1.5)
        ax.axvline(i-0.5, color='black', lw=1.5)
    ax.set_xticks([])
    ax.set_yticks([])
    plt.show()

a = 6
grid = color_grid(a)
plot_color_grid(grid)