Advent of Code 2024 Day 15

from pathlib import Path

TEST = """**REDACTED**"""


def parse(input: str) -> tuple[set[complex], set[complex], complex, list[complex]]:
    walls: set[complex] = set()
    boxes: set[complex] = set()
    robot: complex = 0
    moves: list[complex] = []
    directions = {"<": -1, ">": 1, "^": -1j, "v": 1j}

    for y, row in enumerate(input.splitlines()):
        if row.startswith("#"):
            for x, c in enumerate(row):
                match c:
                    case "#":
                        walls.add(complex(x, y))
                    case "@":
                        robot = complex(x, y)
                    case "O":
                        boxes.add(complex(x, y))
                    case _:
                        pass
        elif row:
            moves.extend([directions[c] for c in row])

    return walls, boxes, robot, moves


def gpscoord(p: complex) -> int:
    return int(100 * p.imag + p.real)


def part1(input: str) -> int:
    walls, boxes, robot, moves = parse(input)
    for d in moves:
        new_robot = robot + d
        if new_robot in walls:
            continue
        if new_robot in boxes:
            new_box = new_robot + d
            while new_box in boxes:
                new_box += d
            if new_box in walls:
                continue
            boxes.add(new_box)
            boxes.remove(new_robot)
        robot = new_robot

    return sum(gpscoord(p) for p in boxes)


assert part1(TEST) == 10092

INPUT = Path("input/day15.txt").read_text()
part1_total = part1(INPUT)
print(f"{part1_total=:,}")  # 1,437,174


def parse2(input: str) -> tuple[set[complex], set[complex], complex, list[complex]]:
    walls: set[complex] = set()
    boxes: set[complex] = set()
    robot: complex = 0
    moves: list[complex] = []
    directions = {"<": -1, ">": 1, "^": -1j, "v": 1j}

    for y, row in enumerate(input.splitlines()):
        if row.startswith("#"):
            for x, c in enumerate(row):
                match c:
                    case "#":
                        walls.add(complex(2 * x, y))
                        walls.add(complex(2 * x + 1, y))
                    case "@":
                        robot = complex(2 * x, y)
                    case "O":
                        boxes.add(complex(2 * x, y))
                    case _:
                        pass
        elif row:
            moves.extend([directions[c] for c in row])

    return walls, boxes, robot, moves


def push_updown(
    direction: complex, box: complex, walls: set[complex], boxes: set[complex]
) -> tuple[bool, set[complex], set[complex]]:
    can_move = True
    nb = box + direction
    add_boxes: set[complex] = {nb}
    remove_boxes: set[complex] = {box}

    if nb in walls or nb+1 in walls:
        return False, set(), set()
    if nb in boxes:
        c, a, r = push_updown(direction, nb, walls, boxes)
        can_move = can_move and c
        add_boxes |= a
        remove_boxes |= r
        return can_move, add_boxes, remove_boxes
    if nb - 1 in boxes:
        c, a, r = push_updown(direction, nb - 1, walls, boxes)
        can_move = can_move and c
        add_boxes |= a
        remove_boxes |= r
    if can_move and nb + 1 in boxes:
        c, a, r = push_updown(direction, nb + 1, walls, boxes)
        can_move = can_move and c
        add_boxes |= a
        remove_boxes |= r
        return can_move, add_boxes, remove_boxes

    return can_move, add_boxes, remove_boxes


def push_right(
    box: complex, walls: set[complex], boxes: set[complex]
) -> tuple[bool, set[complex], set[complex]]:
    nb = box + 1
    if nb in walls or nb + 1 in walls:
        return False, set(), set()
    add_boxes: set[complex] = {nb}
    remove_boxes: set[complex] = {box}
    while nb + 1 in boxes:
        if nb + 3 in walls:
            return False, set(), set()
        add_boxes.add(nb + 2)
        remove_boxes.add(nb + 1)
        nb += 2
    return True, add_boxes, remove_boxes


def push_left(
    box: complex, walls: set[complex], boxes: set[complex]
) -> tuple[bool, set[complex], set[complex]]:
    nb = box - 1
    if nb in walls:
        return False, set(), set()
    add_boxes: set[complex] = {nb}
    remove_boxes: set[complex] = {box}
    while nb - 1 in boxes:
        nb -= 1
        if nb - 1 in walls:
            return False, set(), set()
        add_boxes.add(nb - 1)
        remove_boxes.add(nb)
        nb -= 1

    return True, add_boxes, remove_boxes


def show(walls: set[complex], boxes: set[complex], robot: complex) -> None:
    height = max(int(p.imag) for p in walls) + 1
    width = max(int(p.real) for p in walls) + 1
    for line in range(height):
        points = []
        for x in range(width):
            p = complex(x, line)
            points.append(
                "#"
                if p in walls
                else "["
                if p in boxes
                else "]"
                if p - 1 in boxes
                else "@"
                if p == robot
                else "."
            )
        print("".join(points))


def part2(inp: str, debug:bool=False) -> int:
    walls, boxes, robot, moves = parse2(inp)
    nboxes = len(boxes)
    for d in moves:
        assert nboxes == len(boxes)
        if debug: show(walls, boxes, robot)
        if robot in boxes or robot-1 in boxes:
            print("error!")
            input()

        if debug:
            print("\nMove:", {1:'>', -1: '<', -1j:'^', 1j: 'v'}[d])

        new_robot = robot + d
        if new_robot in walls:
            continue
        c, a, r = True, set(), set()
        if d.real == 0:
            box = new_robot
            if box in boxes:
                c, a, r = push_updown(d, box, walls, boxes)
            elif box - 1 in boxes:
                c, a, r = push_updown(d, box - 1, walls, boxes)
        elif d == 1:
            # push right
            if new_robot in boxes:
                c, a, r = push_right(new_robot, walls, boxes)
        else:  # push left
            if new_robot - 1 in boxes:
                c, a, r = push_left(new_robot - 1, walls, boxes)
        if c:
            robot = new_robot
            boxes = (boxes - r) | a

    show(walls, boxes, robot)
    return sum(gpscoord(p) for p in boxes)


# TEST2 = """**REDACTED**"""
# part2(TEST2, debug=True)
# exit()
assert part2(TEST) == 9021
print(part2(INPUT)) # 1437468