Advent of Code 2024 day 18 part 2

from dataclasses import dataclass
from collections import defaultdict
from math import inf
from argparse import ArgumentParser, FileType


@dataclass(frozen=True, slots=True)
class Coord:
    x: int
    y: int
    z: int = 0

    def taxicab_distance(self, other):
        return abs(self.x - other.x) + abs(self.y - other.y) \
            + abs(self.z - other.z)

    def to_3d(self, z):
        return Coord(self.x, self.y, self.z)

    def on_grid(self, max_x, max_y):
        return 0 <= self.y <= max_x and 0 <= self.x <= max_y

    def move(self, dx=0, dy=0, dz=0):
        return Coord(self.x + dx, self.y + dy, self.z + dz)

    @property
    def neighbours(self):
        yield Coord(self.x, self.y - 1)
        yield Coord(self.x + 1, self.y)
        yield Coord(self.x, self.y + 1)
        yield Coord(self.x - 1, self.y)


def parse_byte(line):
    x, y = line.split(',')
    return Coord(int(x), int(y))


def parse(src):
    return [parse_byte(line) for line in src.splitlines()]


def print_grid(walls, path, max_x, max_y):
    output = ''
    for y in range(max_y+1):
        for x in range(max_x+1):
            coord = Coord(x, y)
            if coord in path:
                output += 'O'
            elif coord in walls:
                output += '#'
            else:
                output += '.'
        output += '\n'
    print(output)


def get_nodes(prev, end):
    nodes = set()
    next_node = end
    while next_node is not None:
        nodes.add(next_node)
        next_node = prev.get(next_node)
    return nodes


def ScoreMap(mapping):
    return defaultdict(lambda: inf, mapping)


def a_star(walls, max_x, max_y):
    start = Coord(0, 0)
    end = Coord(max_x, max_y)

    discovered = set((start,))
    g_score = ScoreMap({start: 0})
    f_score = ScoreMap({start: start.taxicab_distance(end)})
    prev = dict()

    while len(discovered) > 0:
        current_node = min(discovered, key=lambda x: f_score[x])

        if current_node == end:
            return get_nodes(prev, end)

        discovered.remove(current_node)

        for neighbour in current_node.neighbours:
            if not neighbour.on_grid(max_x, max_y) or neighbour in walls:
                continue

            candidate_score = g_score[current_node] + 1
            if candidate_score < g_score[neighbour]:
                g_score[neighbour] = candidate_score
                f_score[neighbour] = candidate_score \
                    + neighbour.taxicab_distance(end)

                discovered.add(neighbour)
                prev[neighbour] = current_node

    return None  # unexpanded nodes exhausted but goal not reached


def main(walls, max_x, max_y, wall_count, verbose):
    # we know from part 1 that there is still a path after some bytes dropped
    end = len(walls)
    while wall_count < end:
        if verbose:
            print(wall_count)
        wall_set = set(walls[:wall_count+1])
        path = a_star(wall_set, max_x, max_y)
        if path is None:
            break

        if verbose:
            print_grid(wall_set, path, max_x, max_y)

        # skip walls that aren't on the path (they don't change anything)
        while wall_count < end and walls[wall_count] not in path:
            wall_count += 1

    final_wall = walls[wall_count]
    return f'{final_wall.x},{final_wall.y}'


arg_parser = ArgumentParser()
arg_parser.add_argument('src', type=FileType('r'))
arg_parser.add_argument('max_x', nargs='?', type=int, default=70)
arg_parser.add_argument('max_y', nargs='?', type=int, default=70)
arg_parser.add_argument('safe_bytes', nargs='?', type=int, default=1024)
arg_parser.add_argument('-v', '--verbose', action='store_true')


if __name__ == '__main__':
    args = arg_parser.parse_args()
    walls = parse(args.src.read())
    print(main(walls, args.max_x, args.max_y, args.safe_bytes, args.verbose))