sudoku generator

sudoku.py

import copy, random
n = 9
values = list(range(n))

def to_matrix(d):
    return [[d[(x,y)] for y in range(n)] for x in range(n)]

def sudoku():
    known = dict()
    unknown = {(x,y): set(values)
               for x in range(n)
               for y in range(n)}

    def assign(x, y, v):
        assert (x,y) not in known
        assert (x,y) in unknown
        known[(x,y)] = v
        del unknown[(x,y)]
        # Remove incompatible choices.
        for (x2,y2), vs in list(unknown.items()):
            assert not (x == x2 and y == y2)
            if x == x2 or y == y2 or (x // 3 == x2 // 3 and y // 3 == y2 // 3):
            # if (x == x2 or y == y2): # don't use diagonals
                vs.discard(v)
                if not vs:
                    # print(f"contradiction at {x2}, {y2}:")
                    # print(f"    known: {known}")
                    # print(f"  unknown: {unknown}")
                    return False
        return True

    def solve():
        nonlocal known, unknown
        # print(f"  known: {known}")
        # print(f"unknown: {unknown}")

        # Have we reached a solution?
        if not unknown:
            yield known
            return

        # Pick a choice with the smallest number of options.
        k,vs = min(unknown.items(), key=lambda x: len(x[1]))

        # If the choice is forced, make it.
        if len(vs) == 1:
            v = vs.pop()
            if not assign(*k, v): return # contradiction!
            yield from solve()
            return

        # Otherwise, backtrack.
        for v in vs:
            # print(f"Unforced choice: {k} -> {v}")
            known_copy = copy.deepcopy(known)
            unknown_copy = copy.deepcopy(unknown)
            if not assign(*k, v): return # contradiction, die!
            yield from solve()
            known = known_copy
            unknown = unknown_copy

    yield from solve()

def main():
    for x in sudoku():
        try: print(to_matrix(x))
        except BrokenPipeError: break

if __name__ == "__main__": main()