Day 17

day17.py

from heapq import heappush, heappop
from dataclasses import dataclass


@dataclass
class Node:
    x: int
    y: int
    score: int
    direction: int
    dir_count: int

    def get_hash(self):
        return hash((self.x, self.y, self.direction % 4, self.dir_count))

    def __lt__(self, other):
        return (self.score < other.score)


with open('input/17.txt', 'r') as f:
    game = [line.strip() for line in f]


def add_node(heap, x, y, score, direction, dir_count):
    directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
    dx, dy = directions[direction % 4]
    x += dx
    y += dy

    if x < 0 or x >= len(game):
        return
    if y < 0 or y >= len(game[0]):
        return

    heappush(heap, Node(x, y, score + int(game[x][y]), direction, dir_count))


end = (len(game) - 1, len(game[0]) - 1)
heap = []
visited = {}
heappush(heap, Node(0, 0, 0, 0, 0))

while heap:
    node = heappop(heap)

    if (node.x, node.y) == end and node.dir_count >= 4:
        print(node)
        break

    if node.get_hash() in visited and visited[node.get_hash()] <= node.score:
        continue
    visited[node.get_hash()] = node.score

    if node.dir_count < 10:
        add_node(heap, node.x, node.y, node.score, node.direction,
                 node.dir_count + 1)

    if node.dir_count >= 4 or (node.x, node.y) == (0, 0):
        add_node(heap, node.x, node.y, node.score, node.direction + 1, 1)
        add_node(heap, node.x, node.y, node.score, node.direction - 1, 1)