A cube puzzle resolution script

cube_solver.py

import numpy as np

segments_lengths = [3, 2, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 1, 2, 1, 1, 2]
assert sum(segments_lengths) == 27

def solve(cube, position, direction, segment_number):
    if segment_number == len(segments_lengths):
        return cube
    x, y, z = direction
    a1, a2 = np.array([y, z, x]), np.array([z, x, y])
    perpendicular_directions = [a1, -a1, a2, -a2]
    for new_direction in perpendicular_directions:
        new_cube, new_position = +cube, +position  # copies
        for _ in range(segments_lengths[segment_number]):
            new_position = new_position + new_direction
            x, y, z = new_position
            in_cube = all([0 <= e < 3 for e in new_position])
            if not in_cube or new_cube[x, y, z]:
                break
            new_cube[x, y, z] = segment_number + 1
        else:
            new_cube = solve(
                new_cube, new_position, new_direction, segment_number + 1
            )
            if new_cube is not None:
                return new_cube
    return None


final_cube = solve(
    cube=np.zeros((3, 3, 3), dtype="int"),
    position=np.array([0, 0, -1]),
    direction=np.array([1, 0, 0]),
    segment_number=0,
)
print(final_cube)