Advent of Code 2024 day 18 part 1

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, byte_count):
    return {parse_byte(next(src)) for _ in range(byte_count)}


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 ScoreMap(mapping):
    return defaultdict(lambda: inf, mapping)


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 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, verbose):
    path = a_star(walls, max_x, max_y)
    if verbose:
        print_grid(walls, path, max_x, max_y)
    return len(path) - 1


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('byte_count', 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, args.byte_count)
    print(main(walls, args.max_x, args.max_y, args.verbose))