Simple Sudoku solver using backtracking

sudoku.py

import numpy as np
import timeit

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

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


def valid(_A, box=True):
    """
    Check if _A is a valid Sudoku solution.
    """
    values = list(range(1, _A.shape[0]+1))
    size = _A.shape[0]
    box_size = 3

    def _row(__A):
        for i in range(size):
            for k in values:
                if k not in __A[i, :]:
                    return False
        return True

    if not _row(_A) or not _row(_A.T):
        return False

    if box:
        step = size // box_size

        for i in range(0, box_size, step):
            for j in range(0, box_size, step):
                for k in values:
                    if k not in _A[i:i+box_size, j:j+box_size]:
                        return False

    return True


def next_ij(_A, i, j):
    if j == _A.shape[0]-1:
        i += 1
        j = 0
    else:
        j += 1

    return i, j


def islastcell(_A, i, j):
    if i == _A.shape[0]-1 and j == _A.shape[0]-1:
        return True
    return False


def check_box(i, j, k, _A):
    ibox = (i // 3) * 3
    jbox = (j // 3) * 3

    if k not in _A[ibox:ibox+3, jbox:jbox+3]:
        return True

    return False


def rowcolbox_valid(i, j, k, _A):
    if k not in _A[i, :] and k not in _A[:, j] and check_box(i, j, k, _A):
        return True
    return False


def solve_backtracking(_A, i, j):
    values = range(1, _A.shape[0]+1)
    i_next, j_next = next_ij(_A, i, j)

    if islastcell(_A, i, j):
        if A[i, j] != 0:
            return valid(_A)

        for k in values:
            _A[i, j] = k

            if valid(_A):
                return True

        A[i, j] = 0
        return False

    if A[i, j] == 0:
        for k in values:
            if rowcolbox_valid(i, j, k, _A):
                _A[i, j] = k

                if solve_backtracking(_A, i_next, j_next):
                    return True

        A[i, j] = 0
        return False

    return solve_backtracking(_A, i_next, j_next)


print(timeit.timeit(lambda: solve_backtracking(A, 0, 0), number=3))

print(f"Is valid Sudoku: {valid(A)}")
print(f"A == A_solution: {np.array_equal(A, A_solution)}")

# ❯ python sudoku.py
# 13.731124277008348
# Is valid Sudoku: True
# A == A_solution: True