Solver for Sudoku-like puzzles with (more or less) arbitrary rules

sudoku.py

import itertools

def is_unique(seq, symbols):
    r = [el for el in seq if el in symbols]
    return len(r) == len(set(r))


def rule_row_unique(board, symbols):
    return all(is_unique(row, symbols) for row in board)


def rule_col_unique(board, symbols):
    return all(is_unique(col, symbols) for col in zip(*board))

def rule_diag_unique(board, symbols):
    d1 = is_unique((board[i][i] for i in range(len(board))), symbols)
    d2 = is_unique((board[i][-1-i] for i in range(len(board))), symbols)
    return d1 and d2

def rule_subsquare_unique(N):
    def _rule_subsquare_unique(board, symbols):
        valid = True
        for i, j in itertools.product(range(0,len(board), N), repeat=2):
            subsquare = []
            for k, l in itertools.product(range(N), repeat=2):
                subsquare.append(board[i+k][j+l])
            valid = valid and is_unique(subsquare, symbols)
        return valid
    return _rule_subsquare_unique


def _sum_valid(seq, N, mnum, symbols):
    '''helper function for msquare rule'''
    seq = list(x for x in seq if x in symbols)
    s = sum(seq)
    if len(seq) == N:
        return s == mnum
    else:
        return s <= mnum
    

def rule_msquare(N):
    '''returns check-rule if board is a partially filled magic square of size N'''
    mnum = sum(i for i in range(N*N+1))//N
    def rule_square(board, symbols):
        uniq = is_unique((board[i][j] for i, j in itertools.product(range(len(board)),repeat=2)), symbols)
        rsum = all(_sum_valid(row, N, mnum, symbols) for row in board)
        csum = all(_sum_valid(col, N, mnum, symbols) for col in zip(*board))
        dsum = all(_sum_valid(d, N, mnum, symbols) for d in zip(*((board[i][i], board[i][-i-1]) for i in range(N))))
        return uniq and rsum and csum and dsum
    return rule_square


def _copy_board(board):
    return [[tuple(field) for field in row] for row in board]


def _choose_value(board, i, j, val):
    b = _copy_board(board)
    b[i][j] = tuple([val])
    return b

def _hide_options(board):
    return [[x[0] if len(x) == 1 else None for x in row] for row in board]

def _check_valid(board, rules, symbols):
    return all(rule(_hide_options(board), symbols) for rule in rules)

def _solver(board, rules, symbols):
    if not _check_valid(board, rules, symbols):
        yield None
    least_options = (len(symbols), -1, -1)
    for i, j in itertools.product(range(len(board)), repeat=2):
        if len(board[i][j]) > 1:
            possible_vals = tuple(val for val in board[i][j] 
                    if _check_valid(_choose_value(board, i, j, val), rules, symbols))
            if len(possible_vals) == 0:
                return
            board[i][j] = possible_vals
            least_options = min(least_options, (len(possible_vals), i, j))
    if all(all(map(lambda f: len(f) <= 1, row)) for row in board):
        yield _hide_options(board)
    else:
        _, i, j = least_options
        for val in board[i][j]:
            yield from _solver(_choose_value(board, i, j, val), rules, symbols)



def puzzle_solutions(board, rules, symbols, placeholder='_'):
    b = []
    for row in board:
        r = []
        for field in row:
            if field in symbols:
                r.append(tuple((field,)))
            elif field == placeholder:
                r.append(tuple(symbols))
            else:
                r.append(tuple())
        b.append(r)
    yield from _solver(b, rules, symbols)



if __name__ == '__main__':
    # puzzle from DnD
    b = '''
0 1 2 3 4 5
3 _ _ _ _ _
5 _ X X _ _
1 _ X X _ _
2 _ _ _ _ _
4 5 3 2 0 1
'''.strip()
    b = [row.split(' ') for row in b.splitlines()]
    symbols = set('012345')
    rules = [rule_row_unique, rule_col_unique, rule_diag_unique]
    solutions = list(puzzle_solutions(b, rules, symbols, '_'))
    print('DnD Pseududoku has {} solutions.'.format(len(solutions)))
    
    # hard sudoku
    b = '''
_ _ _ _ _ _ 6 1 9
_ 2 3 _ 9 _ _ _ _
_ _ _ 1 4 _ _ 2 _
_ _ 9 _ _ _ _ _ _
7 _ 8 _ _ 3 _ _ _
_ _ _ _ _ 5 3 4 _
_ _ _ _ _ _ 4 6 7
8 3 _ _ _ _ _ _ _
_ _ _ 6 1 2 _ _ _
'''.strip()
    b = [row.split(' ') for row in b.splitlines()]
    symbols = set('123456789')
    rules = [rule_row_unique, rule_col_unique, rule_subsquare_unique(3)]
    solution = next(puzzle_solutions(b, rules, symbols, '_'))
    print('\nSolution to the hard sudoku is:')
    print('\n'.join(' '.join(row) for row in solution))
    
    # magic square
    N = 3
    b = [['_' for _ in range(N)] for _ in range(N)]
    symbols = set(range(1,N*N+1))
    rules = [rule_msquare(N)]
    solution = next(puzzle_solutions(b, rules, symbols, '_'))
    print('\nFirst Magic Square of size {} is:'.format(N))
    print('\n'.join(' '.join('{:>2}'.format(f) for f in row) for row in solution))