Sudoku solver using NumPy

Question

I come from the R language and I gradually switch to Python (Numpy and Pandas). I coded this python code for fun to solve a Sudoku grid. Could you tell me if my code is pythonic enough? For example, did I use the numpy functions correctly. How to make the code more readable?

import numpy as np
from functools import reduce

def solver_python(grid):
    numbers=np.arange(1,10)
    i,j = np.where(grid==0) 
    if (i.size==0):
        return(True,grid)
    else:
        i,j=i[0],j[0]    
        row = grid[i,:] 
        col = grid[:,j]
        sqr = grid[(i//3)*3:(3+(i//3)*3),(j//3)*3:(3+(j//3)*3)].reshape(9)
        values = np.setdiff1d(numbers,reduce(np.union1d,(row,col,sqr)))

        grid_temp = np.copy(grid) 

        for value in values:
            grid_temp[i,j] = value
            test = solver_python(grid_temp)
            if (test[0]):
                return(test)

        return(False,None)

example = np.array([[5,3,0,0,7,0,0,0,0],
                    [6,0,0,1,9,5,0,0,0],
                    [0,9,8,0,0,0,0,6,0],
                    [8,0,0,0,6,0,0,0,3],
                    [4,0,0,8,0,3,0,0,1],
                    [7,0,0,0,2,0,0,0,6],
                    [0,6,0,0,0,0,2,8,0],
                    [0,0,0,4,1,9,0,0,5],
                    [0,0,0,0,8,0,0,7,9]])

Answer 1

Overall this is a really good approach. There is one main thing you can do to improve the implementation of this algorithm. You are copying the array at each level of recursion, but it is unnecessary in this case. Instead try an in-place approach. This improves cache locality and saves the tiniest amount of memory.

Secondly, things like (i//3)*3:(3+(i//3)*3) can be confusing. Extracting out a variable name puts less burden on the reader and optimizer as (i//3)*3 is repeated.

You can also break the function up into smaller functions, but where you draw the line is up to you. Your use of reduce here is extremely readable and elegant way of implementing this algorithm. There are better algorithms for solving Sudoku, but since brute force is so fast, this is a good choice for its simplicity.

Here is an example:

import numpy as np
from functools import reduce


def get_candidates(search_space):
    return np.setdiff1d(np.arange(1, 10), reduce(np.union1d, search_space))


def solve(board):
    missing = get_missing(board)

    if not had_missing(missing):
        return True

    missing_col = get_missing_col(missing)
    missing_row = get_missing_row(missing)

    search_space = (
            get_col(board, missing_col),
            get_row(board, missing_row),
            get_square(board, missing_col, missing_row)
        )

    for candidate in get_candidates(search_space):
        board[missing_row, missing_col] = candidate
        if solve(board):
            return True

    board[missing_row, missing_col] = 0
    return False


def get_col(board, idx):
    return board[:, idx].reshape(9)


def get_row(board, idx):
    return board[idx, :].reshape(9)


def get_square(board, col, row):
    col = col // 3 * 3
    row = row // 3 * 3
    return board[row:row+3, col:col+3].reshape(9)


def get_missing(board):
    return np.where(board == 0)


def had_missing(missing):
    return len(missing[0])


def get_missing_col(missing):
    return missing[1][0]


def get_missing_row(missing):
    return missing[0][0]


board = np.array([
    [5,3,0, 0,7,0, 0,0,0],
    [6,0,0, 1,9,5, 0,0,0],
    [0,9,8, 0,0,0, 0,6,0],

    [8,0,0, 0,6,0, 0,0,3],
    [4,0,0, 8,0,3, 0,0,1],
    [7,0,0, 0,2,0, 0,0,6],

    [0,6,0, 0,0,0, 2,8,0],
    [0,0,0, 4,1,9, 0,0,5],
    [0,0,0, 0,8,0, 0,7,9]
])

solve(board)
print(board)
```