not a sudoku solver

not-sudoku.py

import copy, random
n = 4
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

        # Apply all forced choices.
        forced = [(k,vs) for k,vs in unknown.items() if len(vs) == 1]
        if forced:
            # print(f"Forced choices: {forced}")
            for k,vs in forced:
                if not assign(*k, vs.pop()):
                    return # contradiction, die!
            yield from solve()
            return

        # Pick a choice to make and backtrack over all options.
        assert unknown
        known_copy = copy.deepcopy(known)
        unknown_copy = copy.deepcopy(unknown)
        k, vs = random.choice(list(unknown.items()))
        for v in vs:
            # print(f"Unforced choice: {k} -> {v}")
            known = copy.deepcopy(known_copy)
            unknown = copy.deepcopy(unknown_copy)
            if not assign(*k, v): return # contradiction, die!
            yield from solve()

    yield from solve()

if __name__ == "__main__":
    for x in sudoku():
        print(to_matrix(x))