Advent of Code 2024 day 15 part 2

from enum import Enum
from argparse import ArgumentParser, FileType


class Direction(Enum):
    '''A cardinal direction.'''
    UP = (0, -1)
    LEFT = (1, 0)
    DOWN = (0, 1)
    RIGHT = (-1, 0)

    @classmethod
    def from_char(cls, char):
        '''Construct a direction from its corresponding character.'''
        if char == '^':
            return cls.UP
        elif char == '>':
            return cls.LEFT
        elif char == 'v':
            return cls.DOWN
        elif char == '<':
            return cls.RIGHT
        else:
            raise ValueError(f'invalid direction: "{char}"')

    def move(self, x, y):
        '''Translate the coordinates a distance of 1 in this direction.'''
        dx, dy = self.value
        return x + dx, y + dy


def parse_room(src):
    '''Parse the room portion of the input into the room dict and the starting
       coordinates of the robot.'''
    lines = src.split('\n')
    robot = None
    room = dict()

    for y, line in enumerate(lines):
        x = 0
        for char in line:
            if char == '.':
                pass
            elif char == '@':
                robot = (x, y)
            elif char == '#':  # unpushable tiles store None
                room[(x, y)] = None
                room[(x + 1, y)] = None
            elif char == 'O':  # pushable tiles store their pair
                room[(x, y)] = (x + 1, y)
                room[(x + 1, y)] = (x, y)
            else:
                raise ValueError(f'invalid room item: "{char}"')

            x += 2

    assert robot is not None

    return robot, room


def parse(src):
    '''Parse the input into the room dict, the starting coordinates of the
       robot, and an iterable of Directions.'''
    room, moves = src.split('\n\n', 1)
    return (*parse_room(room),
            map(Direction.from_char, moves.replace('\n', '')))


def print_room(robot, room):
    '''Output a string representation of the room and robot to stdout.'''
    width = max(room.keys(), key=lambda item: item[0])[0] + 1
    height = max(room.keys(), key=lambda item: item[1])[1] + 1

    for y in range(height):
        for x in range(width):
            if (x, y) in room:
                if room[(x, y)] is None:
                    print('#', end='')
                elif room[(x, y)][0] > x:
                    print('[', end='')
                else:
                    print(']', end='')
            elif (x, y) == robot:
                print('@', end='')
            else:
                print('.', end='')
        print()
    print()


def wall(x, y, room):
    '''Return true if there is an immovable object at position x, y.'''
    return (x, y) in room and room[(x, y)] is None


def left_box(x, y, room):
    '''Return true if position x, y contains the left side of a box.'''
    return (x, y) in room and room[(x, y)] is not None and room[(x, y)][0] > x


def get_boxes(nx, ny, direction, room):
    '''Construct the dict of boxes that would be affected if the box at
       position x, y where pushed. None is returned if a wall is encountered,
       indicating that x, y cannot be pushed.'''
    boxes = dict()
    todo = {(nx, ny), }
    while len(todo) > 0:
        x, y = todo.pop()

        if wall(x, y, room):  # there's a wall so the boxes can't be moved
            return None

        if (x, y) not in room:
            continue

        pair = room[(x, y)]
        boxes[(x, y)] = pair

        if pair not in boxes:
            todo.add(pair)

        todo.add(direction.move(x, y))

    return boxes


def move(robot, room, direction):
    '''Move the robot in direction, potentially pushing some boxes and updating
       the room.'''
    nx, ny = direction.move(*robot)

    if (nx, ny) not in room:  # free to move
        return (nx, ny)

    if not room[(nx, ny)]:  # blockage, can't do anything
        return robot

    boxes = get_boxes(nx, ny, direction, room)
    if boxes is not None:  # if the next space is free, move them
        robot = direction.move(*robot)
        new_boxes = {direction.move(*box): direction.move(*pair)
                     for box, pair in boxes.items()}

        for old_box in boxes:
            del room[old_box]

        room |= new_boxes

    return robot


def gps_coords(x, y):
    '''Convert x, y to gps coordinates.'''
    return x + 100 * y


def main(robot, room, moves, verbose=False):
    '''Perform each movement and return the sum of the gps coordinates of each
       box.'''
    if verbose:
        print_room(robot, room)

    for direction in moves:
        robot = move(robot, room, direction)

        if verbose:
            print_room(robot, room)

    return sum(gps_coords(*xy) for xy in room if left_box(*xy, room))


arg_parser = ArgumentParser()
arg_parser.add_argument('src', type=FileType('r'))
arg_parser.add_argument('-v', '--verbose', action='store_true')

if __name__ == '__main__':
    args = arg_parser.parse_args()
    print(main(*parse(args.src.read()), args.verbose))