sol.py (1860B) [raw]
1 #!/usr/bin/env python3 2 import sys 3 import heapq 4 from collections import defaultdict 5 6 def parse(): 7 G = [] 8 for l in sys.stdin: 9 G.append([c for c in l.strip()]) 10 # all examples the same w S at bottom left facing east 11 return G, (len(G) - 2, 1), '>' 12 13 def pprint(G): 14 for i in range(len(G)): 15 for j in range(len(G[0])): 16 print(G[i][j], end="") 17 print("") 18 19 def step(pos, dir): 20 i, j = pos 21 if dir == "^": 22 return (i-1, j) 23 elif dir == ">": 24 return (i, j+1) 25 elif dir == "v": 26 return (i+1, j) 27 elif dir == "<": 28 return (i, j-1) 29 else: 30 raise Exception(f"Bad dir: {dir}") 31 32 33 # Return [(cost, newpos, newdir)] 34 def neighs(G, pos, dir): 35 if dir == "^" or dir == "v": 36 return [(1, step(pos, dir), dir), (1001, step(pos, ">"), ">"), (1001, step(pos, "<"), "<")] 37 elif dir == ">" or dir == "<": 38 return [(1, step(pos, dir), dir), (1001, step(pos, "^"), "^"), (1001, step(pos, "v"), "v")] 39 else: 40 raise Exception(f"Bad dir: {dir}") 41 42 def walk(G, pos, dir): 43 h = [] 44 best = 100000 45 dist = defaultdict(lambda: 100000) 46 tiles = set() 47 heapq.heappush(h, (0, pos, dir, [pos])) 48 while h: 49 (cost, pos, dir, pth) = heapq.heappop(h) 50 #print(cost, pos, dir, pth) 51 if cost > dist[(pos, dir)]: 52 continue 53 dist[(pos, dir)] = cost 54 if G[pos[0]][pos[1]] == "E" and cost <= best: 55 best = cost 56 for p in pth: 57 tiles.add(p) 58 for neigh in neighs(G, pos, dir): 59 (ecost, npos, ndir) = neigh 60 if G[npos[0]][npos[1]] == '#': 61 continue 62 heapq.heappush(h, (cost + ecost, npos, ndir, pth + [npos])) 63 return best, len(tiles) 64 65 if __name__ == '__main__': 66 G, pos, dir = parse() 67 print(walk(G, pos , dir))