Aoc day 16 pt 2.py

Aoc day 16 pt2.py

# READ FILE
FILENAME = "A"
f = open(FILENAME).readlines()
f = [x.strip() for x in f]

G=defaultdict()
r=0
P = None
end = None
walls =set()
dist = {}
rot = [(-1,0),(0,1),(1,0),(0,-1)]
for l in f:
    c=0
    for ch in l:
        G[(r,c)] = ch
        if ch == "S":
            P = (r,c)
            for v in rot:
                # Adjacent rotations to EAST (1,0)
                if v == (-1,0) or v == (1,0):
                    dist[(r,c,v)] = 1000
                else:
                    # current state is EAST = (0,1)
                    # so cost is 0
                    if v == (0,1):
                        dist[(r,c,v)] = 0
                    else: # 2 rotations required
                        dist[(r,c,v)] = 2000
        if ch == "E":
            end = (r,c)
            for v in rot:
                dist[(r,c,v)] = float('inf')
        if ch == "#":
            walls.add((r,c))
        if ch == ".":
            for v in rot:
                dist[(r,c,v)] = float('inf')
        c += 1

    r += 1


#Q = deque()
assert(P != None)
ans = 99999999

visited=set()

rot = [(-1,0),(0,1),(1,0),(0,-1)]

best = set()
lowest = {}
visited = {}
seen=set((P[0],P[1],(0,1)))

import heapq
Q = []

prev = defaultdict(set)

Q.append((0,P[0],P[1],(0,1)))

while len(Q) > 0:
    score,r,c,v = heapq.heappop(Q)
    print(score)
    for n in n2((r,c)):
        dv = n[0]-r,n[1]-c
        if n not in walls:
            if dv != v:
                k = rot.index(v)
                cost = 0
                if dv == rot[(k-1)%len(rot)] or dv == rot[(k+1)%len(rot)]:
                    cost = 1*1000
                else:
                    cost = 2*1000
                alt = cost + dist[r,c,v] + 1
                if alt < dist[n[0],n[1],dv]:
                    dist[n[0],n[1],dv] = alt
                    prev[n[0],n[1],dv] = set([(r,c,v)])
                    heapq.heappush(Q, (alt, n[0], n[1], dv))
                elif alt == dist[n[0],n[1],dv]:
                    prev[n[0],n[1],dv].add((r,c,v))
            else:
                alt = 1 + dist[(r,c,v)]
                if alt < dist[n[0],n[1],v]:
                    dist[n[0],n[1],v] = alt
                    prev[n[0],n[1],v] = set([(r,c,v)])
                    heapq.heappush(Q, (alt, n[0], n[1], dv))
                elif alt == dist[n[0],n[1],v]:
                    prev[n[0],n[1],v].add((r,c,v))

# misread question
# you can rotate but just not move forward
# need to be faster

"""
this took me a while because i forgot how dijkstra worked
and had to relearn it. the problems on day 15ish and above require
you to write a bit more code to do them correctly

the idea is to path the graph based on r,c,v state
and get the shortest path 

the trick is to keep track of all the shortest paths with the prev 
state which is a well known (?) dijkstra trick

the trick is if you are about to set the distance to be smaller
override prev

otherwise if the new distance is the same as the existing distance, simply add the source node

since states are not re-expanded since we only add decreasing distances there are no cycles

it is not ok to naively just iterate on states (r,c) because the state space blows up
(r,c,v). as a rule of thumb dist[v] and prev[v] has v always matching your state
"""

best=set()
for v in rot:
    if (end[0],end[1],v) in dist:
        k =dist[(end[0],end[1],v)]
        if k == 82460:
            # run bfs
            Q = deque([(end[0],end[1],v)])
            visited=set()
            while Q:
                r,c, v = Q.popleft()
                best.add((r,c))
                for n in prev[(r,c,v)]:
                    if n not in visited:
                        visited.add(n)
                        Q.append(n)

print("answer", len(best))