Code test solution

solution.py

import sys
import time

"""
Solution Notes:
- The code shown here is larger than necessary.  I think it's easier to discuss the code solution before it's optimized.
  - There are also number of helper functions / features which I could have delete.  
- One of my early design decisions was to go with immutable states. That allows me to playback moves as I spent time learning the rules of the game.
  - This also proved to be an important decision when it came to my own testing and validation. 
- Since I chose to use immutable states I also decided to go with a State Space Search Tree. 
  - https://en.wikipedia.org/wiki/State_space_search
- Each node on in the tree represents a possible state based on an instance of the S3T_node class which has both a parent reference and an array of children nodes.
- To explore the problem I created a Simulator (classes: SimulatorCommand, SimulatorState and Simulator).  Given a position state, the Simulator validates and scores potential moves.
- One of the most important aspects of my solution is the use of an evaluation function to score moves.  I chose to use the number of remaining bikes after a move.  
  I then use the depth in the tree multiplied by 0.01 as a penalty.  So the deeper in the tree a solution is found the lower the score.
  score = child_state.remaining_bikes - ((node.depth + 1) * 0.01)
- I created a Bike_AI class which uses the simulator for move validation and is responsible for building and utilizing the search tree.
  - The Bike_AI uses DFS on the state tree and branch pruning based on the identification of an initial winning move. 
  - This dramatically reduces the size of the resulting state tree and the speed to a solution.
"""
import sys
import time

"""Basic implementation of a State Space Search Tree (S3T)"""
class S3T_Node:
    def __init__(self, data, label=None, parent=None):
        self.data = data
        self.score = 0
        if parent:
            self.depth = parent.depth + 1
        else:
            self.depth = 0
        self.label = f"{label}-{self.depth}"
        self.parent = parent
        self.children = []
    def append_child(self, data):
        self.children.append(data)
    def get_data(self):
        return self.data
    def get_parent(self):
        return self.parent
    def get_depth(self):
        return self.depth
    def get_score(self):
        return self.score
    def get_children(self):
        return self.children
    def set_score(self, score):
        self.score = score

"""
Simulator classes
Classes for building and testing the simulator
"""
class SimulatorCommand:
    COMMAND_NONE = 0
    COMMAND_SPEED = 1
    COMMAND_JUMP = 2
    COMMAND_UP = 3
    COMMAND_DOWN = 4
    COMMAND_SLOW = 5
    COMMAND_WAIT = 6

class SimulatorState:
    def __init__(self, state):
        """constructor"""
        if state == None:
            self.remaining_bikes = 0
            self.speed = 0
            self.bikes = []
            self.command = SimulatorCommand.COMMAND_NONE
            self.game_over = False
            self.last_command = SimulatorCommand.COMMAND_NONE
        else:
            self.remaining_bikes = state.total_bikes
            self.speed = state.speed
            self.bikes = []
            for bike in state.bikes:
                self.bikes.append(bike.copy())
            self.command = self.command
            self.game_over = False
            self.last_command = state.command
    def clone(self):
        new_state = SimulatorState(None)
        new_state.remaining_bikes = self.remaining_bikes
        new_state.speed = self.speed
        new_state.bikes = []
        for bike in self.bikes:
            new_state.bikes.append(bike.copy())
        new_state.command = self.command
        new_state.last_command = self.command
        new_state.game_over = self.game_over
        return new_state

class Simulator:
    """Simulator class. Encapsulates a running instance of a simulation"""

    COMMANDS_STRS = ["", "SPEED", "SLOW", "JUMP", "WAIT", "UP", "DOWN"]
    def __init__(self, simulation_data):
        """constructor"""
        self.name = simulation_data["name"]
        self.total_bikes = simulation_data["total_bikes"]
        self.required = simulation_data["required"]
        self.lanes = [
            list(simulation_data["lanes"][0]),
            list(simulation_data["lanes"][1]),
            list(simulation_data["lanes"][2]),
            list(simulation_data["lanes"][3]),
        ]
        self.max_iterations = len(self.lanes[0])

    def process(self, state):
        """Process a single command or iteration of the simulation"""
        state = state.clone()
        len_bike_data = len(state.bikes)
        is_moving_up = False
        is_moving_down = False

        if state.remaining_bikes == 0:
            state.game_over = True
            return state
        if state.command == SimulatorCommand.COMMAND_SPEED:  # handle SPEED
            state.speed = state.speed + 1
        elif state.command == SimulatorCommand.COMMAND_SLOW:  # handle SLOW
            state.speed = state.speed - 1
            if state.speed < 0:
                state.speed = 0
        elif state.command == SimulatorCommand.COMMAND_JUMP:  # handle JUMP
            pass
        elif state.command == SimulatorCommand.COMMAND_WAIT:  # handle WAIT
            pass
        elif state.command == SimulatorCommand.COMMAND_UP:  # handle UP
            first_active_bike = -1
            for b in range(len_bike_data):
                if state.bikes[b][2]:
                    first_active_bike = b
                    break
            if first_active_bike > -1 and state.bikes[first_active_bike][1] > 0:
                is_moving_up = True
        elif state.command == SimulatorCommand.COMMAND_DOWN:  # handle DOWN
            last_active_bike = -1
            for b in reversed(range(len_bike_data)):
                if state.bikes[b][2]:
                    last_active_bike = b
                    break
            if (
                last_active_bike > -1
                and state.bikes[last_active_bike][1] < len(self.lanes) - 1
            ):
                is_moving_down = True

        for b in range(len_bike_data):
            x, y, a = state.bikes[b]
            if a == 0:  # Bike is inactive, ignore
                continue

            # determine result of speed operation
            landing_slot = x + state.speed
            for j in range(x + 1, landing_slot):
                if j >= self.max_iterations:
                    break

                # if moving up check for holes on current path before moving up
                if (
                    is_moving_up
                    and state.last_command != SimulatorCommand.COMMAND_JUMP
                    and j != landing_slot
                ):
                    if self.lanes[y - 1][j] == "0":
                        state.remaining_bikes = state.remaining_bikes - 1
                        a = 0
                        x = j
                        break

                # if moving down check for holes on current path before moving down
                if (
                    is_moving_down
                    and state.last_command != SimulatorCommand.COMMAND_JUMP
                    and j != landing_slot
                ):
                    if self.lanes[y + 1][j] == "0":
                        state.remaining_bikes = state.remaining_bikes - 1
                        a = 0
                        x = j
                        break

                # if current square has a whole and the last operation was not a jump then lose bike
                if (
                    self.lanes[y][j] == "0"
                    and state.last_command != SimulatorCommand.COMMAND_JUMP
                ):
                    state.remaining_bikes = state.remaining_bikes - 1
                    a = 0
                    x = j
                    break

            if is_moving_up:
                y = y - 1
                state.bikes[b][1] = y

            if is_moving_down:
                y = y + 1
                state.bikes[b][1] = y

            # if land on hole,lose bike
            if (
                landing_slot < self.max_iterations
                and self.lanes[y][landing_slot] == "0"
            ):
                state.remaining_bikes = state.remaining_bikes - 1
                a = 0

            if a != 0:
                x = landing_slot
                if x >= self.max_iterations:
                    x = self.max_iterations
                    state.bikes[b][0] = x - 1
                    state.game_over = True
                else:
                    state.bikes[b][0] = x
            else:
                state.bikes[b][0] = x
                state.bikes[b][2] = 0

        if state.remaining_bikes < self.required:
            state.game_over = True
        if state.game_over == True:
            return state

        if state.remaining_bikes == 0:
            state.game_over = True
            return state

        state.game_over = False
        return state

"""Bike AI module"""
class Bike_AI:
    MAX_DEPTH = 50
    COMMANDS = [
        SimulatorCommand.COMMAND_NONE,
        SimulatorCommand.COMMAND_SPEED,
        SimulatorCommand.COMMAND_JUMP,
        SimulatorCommand.COMMAND_UP,
        SimulatorCommand.COMMAND_DOWN,
        SimulatorCommand.COMMAND_SLOW,
        SimulatorCommand.COMMAND_WAIT,
    ]
    COMMANDS_STRS = [
        "",
        "SPEED",
        "JUMP",
        "UP",
        "DOWN",
        "SLOW",
        "WAIT",
    ]

    def __init__(self, simulator):
        self.start = time.time()
        self.simulator = simulator
        self.total_bikes = simulator.total_bikes
        self.required_bikes = simulator.required
        self.lanes = simulator.lanes
        self.highest_score = -9999
        self.max_depth_reached = 0
        self.winning_node = None
        self.node_id = 0
        self.end_search = False
        self.winning_line = {}

    def process_move(self, data):
        self.node_id = 0
        state = SimulatorState(None)
        state.speed = data["speed"]
        state.remaining_bikes = data["remaining_bikes"]
        for bike in data["bikes"]:
            state.bikes.append(bike.copy())

        self.s3t_node = S3T_Node(state, f"root-{self.node_id}")
        self.build_search_tree(self.s3t_node, 0, self.MAX_DEPTH)
        self.end = time.time()
        return self.get_winning_line()

    def build_search_tree(self, node, depth, max_depth):
        if self.end_search:
            return

        if depth == max_depth or node.get_data().game_over:
            return

        for command in self.COMMANDS:
            if command != 0:
                if node.data.speed == 0 and command != SimulatorCommand.COMMAND_SPEED:
                    continue

                new_state = node.data.clone()
                new_state.command = command

                child_state = self.simulator.process(new_state)
                score = child_state.remaining_bikes - ((node.depth + 1) * 0.01)
                if score < self.highest_score:
                    continue

                self.node_id += 1
                child_node = S3T_Node(child_state, f"{self.node_id}", node)
                child_depth = child_node.get_depth()
                if child_depth > self.max_depth_reached:
                    self.max_depth_reached = child_depth

                child_node.set_score(score)
                node.append_child(child_node)

                game_over = child_state.game_over
                has_win = (
                    game_over and child_state.remaining_bikes >= self.required_bikes
                )
                if has_win:
                    if score > self.highest_score:
                        self.end_search = True
                        self.highest_score = score
                        self.winning_node = child_node

                if not has_win:
                    self.build_search_tree(child_node, depth + 1, max_depth)

    def output_winning_line(self, node):
        if node is None:
            return []
        winning_moves = [node.get_score()]
        while node != None:
            state = node.get_data()
            if state.command == 0:
                break
            winning_moves.append(self.COMMANDS_STRS[state.command])
            node = node.get_parent()
        winning_moves.reverse()
        return winning_moves

    def get_winning_line(self):
        return self.output_winning_line(self.winning_node)

def output(statement):
    print(statement)

def debug(statement):
    print(statement, file=sys.stderr, flush=True)

def main():
    m = int(input())  # the amount of motorbikes to control
    v = int(input())  # the minimum amount of motorbikes that must survive
    l0 = input()  # L0 to L3 are lanes of the road. A dot character . represents a safe space, a zero 0 represents a hole in the road.
    l1 = input()
    l2 = input()
    l3 = input()

    sim_data = {
        "name": "test",
        "total_bikes": m,
        "required": v,
        "speed": 0,
        "lanes": [
            list(l0),
            list(l1),
            list(l2),
            list(l3),
        ],
    }
    
    winning_moves = []

    # game loop
    while True:
        s = int(input())  # the motorbikes' speed
        run_data = {"speed": s, "remaining_bikes": 0, "bikes": []}

        remaining_bikes = 0        
        for i in range(m):
            # x: x coordinate of the motorbike
            # y: y coordinate of the motorbike
            # a: indicates whether the motorbike is activated "1" or detroyed "0"
            x, y, a = [int(j) for j in input().split()]
            run_data["bikes"].append([x, y, a])
            if a > 0:
                remaining_bikes += 1

        run_data["remaining_bikes"] = remaining_bikes
        if remaining_bikes == 0:
            debug("Exit: no more active_bikes")
            break

        sim_data["speed"] = s
        sim_data["total_bikes"] = remaining_bikes
        sim = Simulator(sim_data)
        bike_ai = Bike_AI(sim)
        winning_moves = bike_ai.process_move(run_data)
        winning_moves.pop()  # remove last move which is the branch score

        try:
            command = winning_moves.pop(0)
            output(command)
        except:
            debug("Exit: end of commands")
            break

main()