Sudoku checker and solver for CS258

sudoku_solver.py

hard = [[1,0,0,0,0,7,0,9,0],
        [0,3,0,0,2,0,0,0,8],
        [0,0,9,6,0,0,5,0,0],
        [0,0,5,3,0,0,9,0,0],
        [0,1,0,0,8,0,0,0,2],
        [6,0,0,0,0,4,0,0,0],
        [3,0,0,0,0,0,0,1,0],
        [0,4,0,0,0,0,0,0,7],
        [0,0,7,0,0,0,3,0,0]]

def times(A, B):
    return [(i,j) for i in A for j in B]

dims = range(9)
subdims = ([0,1,2],[3,4,5],[6,7,8])
squares  = times(dims, dims)
unitlist = ([times(dims, [c]) for c in dims] +
            [times([r], dims) for r in dims] +
            [times(rs, cs) for rs in subdims for cs in subdims])
units = dict((s, [u for u in unitlist if s in u])
             for s in squares)
peers = dict((s, set(sum(units[s],[]))-set([s]))
             for s in squares)

def copy(board):
    return map(lambda x: x[:], board)

def subgrid(i):
    return range(i/3*3, i/3*3+3)

def check_row(nums):
    ## each number should be between 0 and 9
    counts = [0 for i in range(10)]
    for num in nums:
        counts[num] += 1
        if num > 0 and counts[num] > 1: 
            return False
    return True

def get_constraints(grid, i, j):
    row = grid[i]
    col = [grid[k][j] for k in dims]
    cell = [grid[k][l] for k in subgrid(i) for l in subgrid(j)]
    if check_row(row) and check_row(col) and check_row(cell):
        if grid[i][j] > 0 and grid[i][j] < 10:
            return str(grid[i][j])
        else:
            return ''.join(map(lambda x: str(x), set(range(1, 10)) - set(row) - set(col) - set(cell)))

def check_sudoku(grid):
    if type(grid) is not list or len(grid) != 9: return
    for row in grid:
        if type(row) is not list or len(row) != 9: return
        for num in row:
            if type(num) is not int or num < 0 or num > 9: return
    for i in range(9):
        for j in range(9):
            if get_constraints(grid, i, j) is None:
                return False
    return True

def erase(board, i, j, d):
    if d not in board[i][j]:
        return board
    board[i][j] = board[i][j].replace(d, '')
    if len(board[i][j]) == 0:
        return False # contradiction
    elif len(board[i][j]) == 1:
        d2 = board[i][j] # single option; erase option from peers
        if not all(erase(board, i1, j1, d2) for (i1,j1) in peers[i,j]):
            return False
    for unit in units[i,j]: # number only present once in unit; assign
        dplaces = [c for c in unit if d in board[c[0]][c[1]]]
        if len(dplaces) == 0:
            return False
        elif len(dplaces) == 1:
            i1,j1 = dplaces[0]
            if not assign(board, i1, j1, d):
                return False
    return board

def assign(board, i, j, d):
    if all(erase(board, i, j, d2) for d2 in board[i][j].replace(d, '')): 
        return board
    else: 
        return False

def assign_from(board):
    return [map(lambda x: int(x), row) for row in board]

def solve_partial(board):
    if board is False: return False
    # most constrained
    cell_lengths = [(len(board[c[0]][c[1]]), c) for c in squares]
    if all(lc[0] == 1 for lc in cell_lengths):
        return board
    least, (i, j) = min([lc for lc in cell_lengths if lc[0] > 1], key=lambda x: x[0])
    for d in board[i][j]:
        newboard = solve_partial(assign(copy(board), i, j, d))
        if newboard: return newboard # found partial solution
    return False

def solve_sudoku (grid):
    ###Your code here.
    grid_check = check_sudoku(grid)
    if grid_check == None or grid_check == False: return grid_check
    board = copy(grid)
    for i in range(9):
        for j in range(9): board[i][j] = get_constraints(grid, i, j)
    soln = solve_partial(board)
    if soln:
        grid = assign_from(soln)
        assert check_sudoku(grid) == True 
        assert sum(map(sum, grid)) == 405
        return grid
    else:
        return False

print solve_sudoku(hard)