Draw the Sudoku graph with a circular layout

circular_sudoku_graph.py

"""
This script draws a Sudoku graph with the vertices arranged in a circle.

David Radcliffe
[email protected]
25 December 2019

The sudoku_graph() function below is being considered for a future version of NetworkX.
"""

import matplotlib.pyplot as plt
import networkx as nx
from networkx.exception import NetworkXError


def sudoku_graph(n=3):
    """Returns the n-Sudoku graph. The default value of n is 3.

    The n-Sudoku graph is a graph with n^4 vertices, corresponding to the
    cells of an n^2 by n^2 grid. Two distinct vertices are adjacent if and
    only if they belong to the same row, column, or n-by-n box.

    Parameters
    ----------
    n: integer
       The order of the Sudoku graph, equal to the square root of the
       number of rows. The default is 3.

    Returns
    -------
    NetworkX graph
        The n-Sudoku graph Sud(n).
    """

    if n < 0:
        raise NetworkXError("The order must be greater than or equal to zero.")

    n2 = n * n
    n3 = n2 * n
    n4 = n3 * n

    # Construct an empty graph with n^4 nodes
    G = nx.empty_graph(n4)

    # The graph has no edges when n = 0 or 1.
    if n < 2:
        return G

    # Add edges for cells in the same row
    for row_no in range(0, n2):
        row_start = row_no * n2
        for j in range(1, n2):
            for i in range(j):
                G.add_edge(row_start + i, row_start + j)

    # Add edges for cells in the same column
    for col_no in range(0, n2):
        for j in range(col_no, n4, n2):
            for i in range(col_no, j, n2):
                G.add_edge(i, j)

    # Add edges for cells in the same box
    for band_no in range(n):
        for stack_no in range(n):
            box_start = n3 * band_no + n * stack_no
            for j in range(1, n2):
                for i in range(j):
                    u = box_start + (i % n) + n2 * (i // n)
                    v = box_start + (j % n) + n2 * (j // n)
                    G.add_edge(u, v)

    return G


if __name__ == '__main__':
    g = sudoku_graph()
    solution = (
        '461523897'
        '938471526'
        '527698341'
        '743152689'
        '256389714'
        '819764253'
        '382916475'
        '174235968'
        '695847132'
    )
    palette = ['red', 'blue', 'green', 'orange', 'purple', 'black', 'brown', 'pink', 'gold']
    colors = [palette[int(digit) - 1] for digit in solution]
    nx.draw_circular(g, node_size=200, node_color=colors)
    plt.gca().set_aspect('equal', adjustable='datalim')
    plt.show()