# Traversing a 2D array with a tree

In order to help me solve the December Jane Street puzzle "Die Agony" I wrote a quick and dirty script to manually brute-force through the thing and do the calculations for me.

I then tried to automate traversing the actual 2D array (a simple list of lists), starting from the first square and going through all possibles tree combinations, until the correct path and score was found, using a recursive function, and my brain melted.

Are there any examples of how this would be done?

Thanks.

# Helper for solving the Jane Street "Die Agony" puzzle,
# with very basic user input validation.

class bcolors:
WARNING = '\033[93m'
ENDC = '\033[0m'

class Die():
'''
Rotatable six-sided die, represented as an array of 6 numbers.
Indices are:
0
123
4
5
Die starts with face 2 facing up and face 1 facing left.
'''
def __init__(self):
# Define die configuration after the four possibles rotations
# (up, down, left, right).
# These are indices for the array below.
self.dir = {
'u': (2, 1, 4, 3, 5, 0),
'd': (5, 1, 0, 3, 2, 4),
'l': (0, 2, 3, 5, 4, 1),
'r': (0, 5, 1, 2, 4, 3)
}

def reset(self):
# Reset the die to unknown values
self.faces = [None, None, None, None, None, None]
self.history = []

def rotate(self, dir: str) -> None:
# Temporary variable to avoid modifying the list in place.
faces = []

# Assign the new faces after rotation
for i in range(6):
faces.append(self.faces[self.dir[dir][i]])
self.faces = faces
self.history.append(self.faces)

def get_face(self):
# Index 2 is the one facing up
return self.faces

def set_face(self, value):
# Set the face up value (index 2) to the number passed
self.faces = value

class Board():
def __init__(self):
self.die = Die()
self.reset()
self.failed = []

def reset(self):
self.n = 0
self.score = 0
self.previous = 0
self.steps = []
self.path = 
self.die.reset()

def undo(self):
# To undo a step, we recreate all previous steps minus the last one.
# This was easier than implementing a proper undo stack.
history = self.steps
path = self.path
self.reset()
for i, step in enumerate(history[0:-1]):
self.move([step, path[i+1]])

def move(self, cmd):
# Get the direction for the rotation
# and the goal number of the square the die rotates into.
dir = cmd
goal = int(cmd)

self.die.rotate(dir)

# If the side facing up isn't known,
# set it to the value that matches the goal.
if not self.die.get_face():
self.die.set_face((goal - self.path[self.n]) / (self.n+1))
print(f'\tSetting face to: {self.die.get_face()} (goal: {goal}).')

# Update path, history and score
self.steps.append(dir)
self.n += 1
self.score += self.n * self.die.get_face()
self.path.append(int(self.score))

# Undo if the score doesn't match the goal
# (this will only happen if an unknown face hasn't been set above).
if goal != self.score:
print(f'{bcolors.WARNING}\t\t\t\t\t*** Wrong score ***{bcolors.ENDC}')
self.fail(goal)
self.undo()

# Undo if the set face is not an integer.
elif self.die.get_face() != int(self.die.get_face()):
print(f'{bcolors.WARNING}\t\t\t\t*** No integer face *** {bcolors.ENDC}')
self.fail(goal)
self.undo()

def fail(self, goal):
# Add the failed path to the failed attempts history list
self.path[-1] = goal
print(f'{bcolors.WARNING}\t\t\t\t\t{self.path}{bcolors.ENDC}')
self.failed.append(self.path)

def show_failed(self):
# Display the failed attemps history list
for item in self.failed:
print(item)

def print(self):
# Print the current step, path, die and score
print(f'\nSTEP {self.n}')
print(f'{self.steps} {self.path}')
print(f'Face {self.die.get_face()}')
print(f'Score {self.score}')
print(f'Die {self.die.faces}')

board = Board()

while True:
board.print()

# Get user input
cmd = input('\n> ').split(' ')

if cmd == 'b':   # Go back a step
board.undo()
if cmd == 'f':   # Show failed attempts history
board.show_failed()
if cmd == 'q':   # Quit
break
elif len(cmd) == 2 and cmd in 'udlr' and -7 <= int(cmd) <= 732:
board.move(cmd) # Do rotation with min and max goal value of puzzle