Another sudoku solver

sudoku_solver.py

from copy import deepcopy
from pprint import pprint
import numpy as np

# Idea borrowed from the video https://youtu.be/G_UYXzGuqvM
# and implemented in a kind of numpyish way
#
# TODO: 1) how can I use numpy.lib.stride_tricks.sliding_window_view 
#          or something like that to solve the task?
#       2) get rid of the array SOLUTION
#
#           SUDOKU 9×9 MATRIX
#              as a 4D-Cube
#
#                    k
#          ╭───────────────────╮
#            l       l       l
#          ╭───╮   ╭───╮   ╭───╮
#          _ _ _   _ _ _   _ _ _
#  ╭   ╭ |       |       |       |   
#  │  j│ |  3×3  |       |       |   
#  │   ╰ | _ _ _ | _ _ _ | _ _ _ |   
#  │   ╭ |       |       |       |   
# i│  j│ |       |       |       |   
#  │   ╰ | _ _ _ | _ _ _ | _ _ _ |   
#  │   ╭ |       |       |       |   
#  │  j│ |       |       |       |   
#  ╰   ╰ | _ _ _ | _ _ _ | _ _ _ |   
#
#  where i, j, k, l ∈ {0, 1, 2}

CUBE_WIDTH = 3
ROWS = COLS = 9
NUMBER_OF_CELLS = ROWS*COLS
DIGITS = set(range(1, ROWS+1))     # a set of possible values in cells: 1, 2, ... 9

SUDOKU = np.asarray([              # just one of sudoku puzzles (with 2 solutions)
    [0,2,0, 0,1,0, 0,0,0],         # zero means the cell is empty
    [0,5,1, 0,0,0, 0,8,6],
    [0,0,0, 0,0,9, 2,0,0],
    
    [0,0,6, 7,9,0, 0,0,0],
    [7,8,0, 1,4,0, 6,9,3],
    [5,4,9, 2,3,0, 1,0,8],
    
    [0,0,5, 0,8,0, 0,0,0],
    [4,0,0, 0,0,1, 0,2,9],
    [0,0,0, 0,6,4, 8,0,7]
], dtype=np.uint8).reshape([CUBE_WIDTH]*4)

SOLUTION = deepcopy(SUDOKU)


def possible_values(i, j, k, l) -> set:
    '''Return possible digits for the cell SOLUTION[i,j,k,l]''' 
    if SUDOKU[i, j, k, l] != 0:
        return {SUDOKU[i, j, k, l]}
    else:
        return DIGITS - {*SOLUTION[i, j, :, :].flat,    # digits in a row [i,j]
                         *SOLUTION[:, :, k, l].flat,    # digits in a column [k,l]
                         *SOLUTION[i, :, k, :].flat}    # digits in a 3x3-square [i,k]
                         
                         
def get_solution(flat_index=0) -> np.ndarray:
    '''Get a solution of a SUDOKU puzzle'''
    if flat_index == NUMBER_OF_CELLS:
        yield SOLUTION.reshape(ROWS, COLS)    # TODO: deepcopy?
    else:
        # row, column = flat_index // 9, flat_index % 9 
        # i, j = row // 3, row % 3
        # k, l = column // 3, column % 3
        # here ROWS == COLUMNS == 9, 
        # inner square width aka cube width == 3
        row, col = divmod(flat_index, ROWS)
        i, j, k, l = *divmod(row, CUBE_WIDTH), *divmod(col, CUBE_WIDTH)
        for value in possible_values(i, j, k, l):
            SOLUTION[i, j, k, l] = value
            yield from get_solution(flat_index+1)
        # restore the current cell before returning back
        SOLUTION[i, j, k, l] = SUDOKU[i, j, k, l]
        
        
for solution in get_solution():
    pprint(solution)