snake_cube_code

puzzlesolver.py

'''
    Allen Ma 2019
    Accompanying blog post: https://allenma.me/posts/snakepuzzle/
'''

import numpy as np


# length of the cube dimension - for a 4x4x4 cube it's 4 for instance
LEN_DIM = 3


TOTAL_CUBES = LEN_DIM ** 3


'''
    main algorithm - uses backtracking with a stack

    pass in a list of positions representing the current state of the puzzle
    each position is a 3 element np array (x, y, z)

    returns move stack, a list of moves to make
    (4, 270) for instance means rotate index 4 (hold cube 4 still - cubes is 0 indexed) and rotate the preceeding cubes (0, 1, 2, 3) by 270 counterclockwise
'''
def solve(cubes):
    cur_idx = 1
    cur_angle_idx = 0
    move_stack = []
    angles = [0, 90, 180, 270]
    while not solved(cur_idx):
        # for the current index, try each angle in turn
        for i in range(cur_angle_idx, len(angles)):
            angle = angles[i]
            res, newstate = rotate_puzzle(cur_idx, angle, cubes)
            # check legality of move
            if res and not out_of_bounds(cur_idx, cubes):
                cubes = [c[:] for c in newstate]
                
                # add the move onto the stack, and try rotating the next index
                move_stack.append([cur_idx, i])
                cur_idx += 1
                cur_angle_idx = 0
                break
        # the move was not legal, so we have to redo
        else:
            idx, angle_idx = move_stack.pop()
            cubes = undo_move(idx, angles[angle_idx], cubes)
            while angle_idx == len(angles) - 1:
                idx, angle_idx = move_stack.pop()
                cubes = undo_move(idx, angles[angle_idx], cubes)
            
                if idx == 0:
                    print("there is no solution to this problem")
                    return

            cur_idx = idx
            cur_angle_idx = angle_idx + 1
            assert(cur_angle_idx <= len(angles) - 1)


    print("wooohooo!!! solved - here is the sequence of moves starting from the initial configuration")
    for p in move_stack:
        print(f"index: {p[0]}, angle: {angles[p[1]]}")
    
    mvs = [(m[0], angles[m[1]]) for m in move_stack]
    return mvs



'''
    takes in list of tuples, each tuple representing (idx_rotation, angle)
    returns list of lists, not list of np arrays representing cube state
'''
def perform_rotations(rotations):
    cubes = reset_cubes()
    for r in rotations:
        ok, cubes = rotate_puzzle(r[0], r[1], cubes)
        if not ok:
            print(f"rotation {r[0]}, {r[1]} is invalid")
            return False
    cube_list = [c.tolist() for c in cubes]
    return cube_list

'''
    Checks if the preceeding cubes preceeding cubes[idx] are more than LEN_DIM - 1 away by any dimension
    Returns true if out of bounds, false otherwise.
    For example, for a 3x3x3 cube, if idx==4, then the function would check if any of cubes 0, 1, 2, 3 had a vector difference
    between itself and cube[4] or greater than 2 in any dimension, since that would imply that a cube is out of bounds
'''
def out_of_bounds(idx, cubes):
    # for the special 3x3 case, the maximum absolute value of any component is 2
    MAX_COMP_DIST = LEN_DIM - 1
    # translate the pivot cube to the center, and check if any preceeding cubes are out
    
    # make a copy of the cubes positions
    cubes_copy = [c[:] for c in cubes]
    # print("checking for out of bounds")
    for i in range(idx + 1):
        i_pos = cubes_copy[i]
        i_pos = i_pos - cubes_copy[idx]
        # print(i_pos)
        # check if a cube is out of bounds
        if (abs(i_pos[0]) > MAX_COMP_DIST) or (abs(i_pos[1]) > MAX_COMP_DIST) or (abs(i_pos[2]) > MAX_COMP_DIST):
            return True

    return False

'''
    Undoes rotate_puzzle - takes in cube idx, angle, and current state
    Undoing 90 CCW is rotating 270 CCW for instance
    Returns state after undoing
'''
def undo_move(idx, angle, cubes):
    res = False
    if angle == 0:
        res = True
    elif angle == 90:
        res, cubes = rotate_puzzle(idx, 270, cubes)
    elif angle == 180:
        res, cubes = rotate_puzzle(idx, 180, cubes)
    elif angle == 270:
        res, cubes = rotate_puzzle(idx, 90, cubes)
    else:
        raise ValueError(f"{angle_idx} not valid")

    if not res:
        raise ValueError("undoing not valid for some reason...")
    
    return cubes
        

'''
    checks for a win condition - basically for a 3x3x3 cube, if the current idx is 26, then at least 26 cubes must be in the correct configuration, so the
    cube must be solved
    Returns true for solved
'''
def solved(idx):
    return idx == TOTAL_CUBES - 1

'''
    takes in a cube with which to rotate along the only possible axis of the cube - we'll call this the pivot cube
    def - either 90, 180 or 270 (360 is redundant) - nb rotation by 90 clockwise == rotation by 270 counterclockwise, thus we only
        need to consider one direction
    returns a tuple:
        - false if no rotation is possible, returns the state passed in
        - true otherwise, and returns the new state
'''
def rotate_puzzle(idx, deg, cubes):
    curcube_pos = cubes[idx]
    if idx == 0:
        print("cannot rotate by the 0th position")
        return False, cubes

    if deg == 0:
        return True, cubes

    dummy_pos = [c[:] for c in cubes]

    # calculate the only possible axis to make the rotation
    axis = abs(cubes[idx - 1] - curcube_pos)
    assert(np.sum(axis) == 1)
    if axis.tolist() == [1, 0, 0]:
        ax = "x"
    elif axis.tolist() == [0, 1, 0]:
        ax = "y"
    elif axis.tolist() == [0, 0, 1]:
        ax = "z"
    else:
        raise ValueError

    # try a rotation, and see if it is legal
    # note that rotations about a cube only affect the cubes AFTER the pivot cube
    for i, dpos in enumerate(dummy_pos[:idx]):
        # calculate the vector to c
        vec = dpos - curcube_pos
        # rotate 
        new_vec = rotate_vec(deg, ax, vec)
        # add the vector back to c
        dummy_pos[i] = curcube_pos + new_vec
    
    # check for any duplicate coordinates, indicating that the rotation failed
    seen = []
    for p in dummy_pos:
        for s in seen:
            if (s == p).all(): 
                # print(f"error: rotation invalid because two values {s}, {p} are identical")
                return False, cubes
        else:
            seen.append(p)
    
    return True, dummy_pos


'''
    utility function for rotating a vector by 90, 180 or 270
    Could have used sin, cos, but we are dealing w. orthogonal rotations
    So the math is a lot easier
    we are always assuming counterclockwise rotation
    Returns np array of output vector rotated
'''
def rotate_vec(deg, axis, vec):

    # break up the vector into ai + bj + ck
    # where i, j, k are the standard unit vectors with positive components
    a, b, c = vec[0], vec[1], vec[2]

    # rotation about x axis
    lookup_x = {
        "i": {
            90: [1, 0, 0],
            180: [1, 0, 0],
            270: [1, 0, 0],
        },
        "j": {
            90: [0, 0, -1],
            180: [0, -1, 0],
            270: [0, 0, 1],
        },
        "k": {
            90: [0, 1, 0],
            180: [0, 0, -1],
            270: [0, -1, 0],
        }
    }

    # rotation about y_axis
    lookup_y = {
        "i": {
            90: [0, 0, 1],
            180: [-1, 0, 0],
            270: [0, 0, -1]
        },
        "j": {
            90: [0, 1, 0],
            180: [0, 1, 0],
            270: [0, 1, 0]
        },
        "k": {
            90: [-1, 0, 0],
            180: [0, 0, -1],
            270: [1, 0, 0]
        }
    }

    # rotation about z axis
    lookup_z = {
        "i": {
            90: [0, 1, 0],
            180: [-1, 0, 0],
            270: [0, -1, 0]
        },
        "j": {
            90: [-1, 0, 0],
            180: [0, -1, 0],
            270: [1, 0, 0]
        },
        "k": {
            90: [0, 0, 1],
            180: [0, 0, 1],
            270: [0, 0, 1]
        }
    }

    if axis == "x":
        return a * np.array(lookup_x["i"][deg]) + b * np.array(lookup_x["j"][deg]) + c * np.array(lookup_x["k"][deg])
    elif axis == "y":
        return a * np.array(lookup_y["i"][deg]) + b * np.array(lookup_y["j"][deg]) + c * np.array(lookup_y["k"][deg])
    elif axis == "z":
        return a * np.array(lookup_z["i"][deg]) + b * np.array(lookup_z["j"][deg]) + c * np.array(lookup_z["k"][deg])
    else:
        print(f"cannot rotate by axis {axis}")

def reset_cubes(cubes = None):

    if cubes != None:
        return cubes

    # for a 2x2x2 cube, for example - change LEN_DIM to 2
    # dirs = ["y", "z", "-y", "-x", "y", "-z", "-y"]

    # for the flat snake going up x-y
    dirs = ["x", "x", "y", "x", "y", "y", "x", "y", "y", "x", "y", "x", "x", "y", "y", "x", "y", "x", "y", "y", "x", "x", "y", "y", "x", "x"]

    # construct the cube
    CUBES = [np.array([0, 0, 0])]
    curCube_pos = CUBES[0]
    for i in range(len(dirs)):
        d = dirs[i]
        if d == "x":
            c = curCube_pos + np.array([1, 0, 0])
        elif d == "-x":
            c = curCube_pos + np.array([-1, 0, 0])
        elif d == "y":
            c = curCube_pos + np.array([0, 1, 0])
        elif d == "-y":
            c = curCube_pos + np.array([0, -1, 0])
        elif d == "z":
            c = curCube_pos + np.array([0, 0, 1])
        elif d == "-z":
            c = curCube_pos + np.array([0, 0, -1])
        else:
            raise ValueError(f"{d} is not valid")
        
        CUBES.append(c)
        curCube_pos = c

    return CUBES

def main():
    cubes = reset_cubes()
    print("initially cubes is")
    print(cubes)
    print("solving - might take like 5 mins")
    solve(cubes)

    

if __name__ == "__main__":
    main()
    

Think I found a bug, see comment on your website:
https://allenma.me/posts/snakepuzzle/