aoc2024/day16.py

#!/usr/bin/env python3
import math
from enum import Enum, auto
from typing import Tuple, Dict, Set
from copy import copy


class Tile(Enum):
    WALL = auto()
    PATH = auto()

class Direction(Enum):
    NORTH = auto()
    EAST = auto()
    SOUTH = auto()
    WEST = auto()

type Pos = Tuple[int,int]

maze = []
end: Pos = (0,0)
player: Pos = (0,0)

with open("day16") as f:
    for y,l in enumerate(f):
        temp = []
        for x,c in enumerate(l):
            if c == '#':
                temp.append(Tile.WALL)
            else:
                temp.append(Tile.PATH)

            if c == 'E':
                end = (y,x)

            if c == 'S':
                player = (y,x)

        maze.append(temp)


def next_moves(pos, direction, cost):
    moves = []

    for d in Direction:
        if((d == Direction.NORTH and direction == Direction.SOUTH)
           or (d == Direction.SOUTH and direction == Direction.NORTH)
           or (d == Direction.WEST and direction == Direction.EAST)
           or (d == Direction.EAST and direction == Direction.WEST)):
            continue # only 90°

        p = pos
        if d == Direction.NORTH:
            p = (p[0]-1, p[1])
        if d == Direction.EAST:
            p = (p[0], p[1]+1)
        if d == Direction.SOUTH:
            p = (p[0]+1, p[1])
        if d == Direction.WEST:
            p = (p[0], p[1]-1)

        # can't move into wall
        if maze[p[0]][p[1]] == Tile.WALL: continue

        # because I'm a good, defensive programmer lol
        if p[0] < 0 or p[1] < 0: raise RuntimeError("Crossed a wall")
        if p[0] >= len(maze) or p[1] >= len(maze[0]): raise RuntimeError("Crossed a wall")


        c = cost
        if d == direction: c += 1
        else: c += 1000+1

        moves.append((p,d,c))

    return moves

visited: Dict[Pos, Set[Tuple[int, Direction]]] = {}
queue = [(player, Direction.EAST, 0)]
hits = []
cheap_map = [[math.inf for _ in range(len(maze[0]))] for _ in range(len(maze))]

while queue:
    w = queue.pop()
    if w[0] in visited:
        v = visited[w[0]]
        v0 = next(iter(visited[w[0]])) # just get any element, for cost doesn't matter
        # this is not the cheapest path, prune search here
        if v0[0] < w[2]: continue
        # this is the cheapest path so far, replace all others found so far
        if v0[0] > w[2]: visited[w[0]] = {(w[2], w[1])}
        # another possibility, equally good as the others, but potentially different direction
        else: visited[w[0]].add((w[2], w[1]))
    else:
        visited[w[0]] = {(w[2], w[1])}

    if w[2] < cheap_map[w[0][0]][w[0][1]]:
        cheap_map[w[0][0]][w[0][1]] = w[2]

    if w[0] == end:
        hits.append(w[2])

    queue.extend(next_moves(w[0], w[1], w[2]))

print(min(hits))

tiles = set()
queue = [(player, Direction.EAST, 0, [player])]
while True:
    if not queue: break

    w = queue.pop()

    if w[0] == end and w[2] == min(hits):
        tiles.update(w[3])
        continue

    tainted = False
    # we might cross a path that is cheaper at this point
    # but which has to take a turn while we don't
    if cheap_map[w[0][0]][w[0][1]] < w[2]:
        tainted = True

    nms = next_moves(w[0], w[1], w[2])

    for nm in nms:
        # if we're tainted we are either back
        # to cheapest path here or we won't come back
        if tainted and cheap_map[nm[0][0]][nm[0][1]] < nm[2]:
            continue
        l = copy(w[3]); l.append(nm[0]);
        queue.append((nm[0], nm[1], nm[2], l))

print(len(tiles))

def pprint():
    for i,y in enumerate(maze):
        for j,x in enumerate(y):
            if (i,j) in tiles:
                print("O", sep=" ", end="")
            elif x == Tile.WALL:
                print("#", sep=" ", end = "")
            elif x == Tile.PATH:
                print(".", sep=" ", end = "")
        print()

def pprint_cheap():
    for y in cheap_map:
        for x in y:
            if x == math.inf:
                print("#"*6, sep=" ", end = "")
            else:
                print(f" {x:<5}", sep=" ", end="")
        print()