|
|
@@ -0,0 +1,136 @@
|
|
|
+#!/usr/bin/env python3
|
|
|
+import sys
|
|
|
+from heapq import heappop, heappush
|
|
|
+from itertools import chain, cycle
|
|
|
+
|
|
|
+def read_grid(f):
|
|
|
+ rows = [line.rstrip() for line in f]
|
|
|
+ return list(chain.from_iterable(rows)), len(rows[0])
|
|
|
+
|
|
|
+def build_graph(grid, w):
|
|
|
+ def dfs(loc, prev, dist):
|
|
|
+ nblocs = (loc - 1, loc + 1, loc - w, loc + w)
|
|
|
+ neighbors = [nb for nb in nblocs if grid[nb] != '#']
|
|
|
+
|
|
|
+ node = grid[loc]
|
|
|
+ if node == '@' or node.isalpha() or len(neighbors) > 2:
|
|
|
+ if node in doubles:
|
|
|
+ seen = doubles[node]
|
|
|
+ count = seen.setdefault(loc, len(seen))
|
|
|
+ node += str(count)
|
|
|
+
|
|
|
+ if prev is None:
|
|
|
+ graph[node] = {}
|
|
|
+ else:
|
|
|
+ graph[prev][node] = dist
|
|
|
+ graph.setdefault(node, {})[prev] = dist
|
|
|
+
|
|
|
+ prev = node
|
|
|
+ dist = 0
|
|
|
+
|
|
|
+ if loc not in visited:
|
|
|
+ visited.add(loc)
|
|
|
+ for nb in neighbors:
|
|
|
+ dfs(nb, prev, dist + 1)
|
|
|
+
|
|
|
+ doubles = {'@': {}, '.': {}}
|
|
|
+ graph = {}
|
|
|
+ visited = set()
|
|
|
+ for i, node in enumerate(grid):
|
|
|
+ if node == '@':
|
|
|
+ dfs(i, None, 0)
|
|
|
+ return graph
|
|
|
+
|
|
|
+def collect_keys_alone(graph):
|
|
|
+ totalkeys = sum(node.islower() for node in graph)
|
|
|
+ inf = 10000000000
|
|
|
+ work = [(0, totalkeys, '@0', '')]
|
|
|
+ best_dists = {('@0', ''): 0}
|
|
|
+
|
|
|
+ def visit(node, dist, keys, remkeys):
|
|
|
+ ident = node, keys
|
|
|
+ if best_dists.get(ident, inf) > dist:
|
|
|
+ best_dists[ident] = dist
|
|
|
+ heappush(work, (dist, remkeys, node, keys))
|
|
|
+
|
|
|
+ while work:
|
|
|
+ dist, remkeys, node, keys = heappop(work)
|
|
|
+ if remkeys == 0:
|
|
|
+ return dist
|
|
|
+
|
|
|
+ for nb, step in graph[node].items():
|
|
|
+ if nb.islower() and nb not in keys:
|
|
|
+ nbkeys = ''.join(sorted(keys + nb))
|
|
|
+ visit(nb, dist + step, nbkeys, remkeys - 1)
|
|
|
+ elif not nb.isupper() or nb.lower() in keys:
|
|
|
+ visit(nb, dist + step, keys, remkeys)
|
|
|
+
|
|
|
+def split_grid(grid, w):
|
|
|
+ e = grid.index('@')
|
|
|
+ grid[e - w - 1:e - w + 2] = '@#@'
|
|
|
+ grid[e - 1:e + 2] = '###'
|
|
|
+ grid[e + w - 1:e + w + 2] = '@#@'
|
|
|
+
|
|
|
+def collect_from(graph, root, remkeys, keys, dist):
|
|
|
+ inf = 100000000
|
|
|
+ work = [(dist, remkeys, root, keys)]
|
|
|
+ best_dists = {(root, keys): dist}
|
|
|
+
|
|
|
+ def visit(node, dist, keys, remkeys):
|
|
|
+ ident = node, keys
|
|
|
+ if best_dists.get(ident, inf) > dist:
|
|
|
+ best_dists[ident] = dist
|
|
|
+ heappush(work, (dist, remkeys, node, keys))
|
|
|
+
|
|
|
+ best = dist, remkeys, root, keys
|
|
|
+
|
|
|
+ while True:
|
|
|
+ while work:
|
|
|
+ dist, remkeys, node, keys = heappop(work)
|
|
|
+ if remkeys == 0:
|
|
|
+ yield True, keys, dist
|
|
|
+ return
|
|
|
+
|
|
|
+ if remkeys < best[1] or (remkeys == best[1] and dist < best[0]):
|
|
|
+ best = dist, remkeys, node, keys
|
|
|
+
|
|
|
+ for nb, step in graph[node].items():
|
|
|
+ if nb.islower() and nb not in keys:
|
|
|
+ nbkeys = ''.join(sorted(keys + nb))
|
|
|
+ visit(nb, dist + step, nbkeys, remkeys - 1)
|
|
|
+ elif not nb.isupper() or nb.lower() in keys:
|
|
|
+ visit(nb, dist + step, keys, remkeys)
|
|
|
+
|
|
|
+ dist, remkeys, node, keys = best
|
|
|
+ newkeys, newdist = yield False, keys, dist
|
|
|
+ newremkeys = remkeys - (len(newkeys) - len(keys))
|
|
|
+ work.append((newdist, newremkeys, node, newkeys))
|
|
|
+
|
|
|
+def collect_keys(graph):
|
|
|
+ entrances = tuple(node for node in graph if node[0] == '@')
|
|
|
+ totalkeys = sum(node.islower() for node in graph)
|
|
|
+ bots = []
|
|
|
+ keys = ''
|
|
|
+ dist = 0
|
|
|
+
|
|
|
+ for entrance in entrances:
|
|
|
+ bot = collect_from(graph, entrance, totalkeys - len(keys), keys, dist)
|
|
|
+ bots.append(bot)
|
|
|
+ done, keys, dist = next(bot)
|
|
|
+ if done:
|
|
|
+ return dist
|
|
|
+
|
|
|
+ for bot in cycle(bots):
|
|
|
+ done, keys, dist = bot.send((keys, dist))
|
|
|
+ if done:
|
|
|
+ return dist
|
|
|
+
|
|
|
+# part 1
|
|
|
+grid, w = read_grid(sys.stdin)
|
|
|
+graph = build_graph(grid, w)
|
|
|
+print(collect_keys(graph))
|
|
|
+
|
|
|
+# part 2
|
|
|
+split_grid(grid, w)
|
|
|
+graph = build_graph(grid, w)
|
|
|
+print(collect_keys(graph))
|
Taddeus Kroes