2019 day 18

2019/d18.py

@@ -0,0 +1,199 @@
+from lib import get_data, Grid2D, LETTERS_UPPER, LETTERS_LOWER, add2
+from collections import defaultdict
+
+LETTERS = LETTERS_UPPER + LETTERS_LOWER
+
+data = """#############
+#g#f.D#..h#l#
+#F###e#E###.#
+#dCba...BcIJ#
+#####.@.#####
+#nK.L...G...#
+#M###N#H###.#
+#o#m..#i#jk.#
+#############"""
+
+
+def part_1(data):
+    g = Grid2D(data)
+    (start,) = g.find("@")
+    starts = [start]
+
+    graph = defaultdict(set)
+    starts_seen = set()
+
+    while starts:
+        start = starts.pop()
+        if start in starts_seen:
+            continue
+        else:
+            starts_seen.add(start)
+
+        start_symbol = g[start]
+        xs = [start]
+        seen = set()
+
+        for steps in range(1_000):
+            nxs = []
+
+            for x in xs:
+                if x in seen:
+                    continue
+                else:
+                    seen.add(x)
+
+                for nb in g.neighbors_ort(x):
+                    if g[nb] == ".":
+                        nxs.append(nb)
+                    elif g[nb] in LETTERS:
+                        symbol = g[nb]
+                        if symbol != start_symbol:
+                            graph[symbol].add((start_symbol, steps + 1))
+                            graph[start_symbol].add((symbol, steps + 1))
+                        starts.append(nb)
+            xs = nxs
+            if len(xs) == 0:
+                break
+
+    all_keys = [g[p] for p in g.find(LETTERS_LOWER)]
+    poss = [(0, 0, tuple("@"), "@")]
+    best: dict[tuple[tuple, str], int] = {(("@",), "@"): 0}
+
+    min_dist = 10**9
+    while poss:
+        current_distance, key_count, keys, symbol = poss.pop()
+        # print(current_distance, key_count, keys, symbol)
+
+        if key_count - 1 == len(all_keys):
+            min_dist = min(min_dist, current_distance)
+            continue
+
+        for next_symbol, distance in graph[symbol]:
+            if next_symbol in LETTERS_UPPER and not next_symbol.lower() in keys:
+                continue
+
+            if next_symbol in LETTERS_LOWER:
+                new_keys = set(keys)
+                new_keys.add(next_symbol)
+                new_keys = tuple(sorted(new_keys))
+                new_key_count = len(new_keys)
+            else:
+                new_keys = keys
+                new_key_count = key_count
+
+            new_distance = current_distance + distance
+            key = (new_keys, next_symbol)
+            if (key not in best) or (key in best and best[key] > new_distance):
+                best[key] = new_distance
+                poss.append((new_distance, new_key_count, new_keys, next_symbol))
+    print(min_dist)
+
+
+def part_2(data):
+    g = Grid2D(data)
+    (start,) = g.find("@")
+    g[start] = "#"
+    g[add2(start, (-1, 0))] = "#"
+    g[add2(start, (1, 0))] = "#"
+    g[add2(start, (0, 1))] = "#"
+    g[add2(start, (0, -1))] = "#"
+
+    g[add2(start, (-1, -1))] = "0"
+    g[add2(start, (-1, 1))] = "1"
+    g[add2(start, (1, 1))] = "2"
+    g[add2(start, (1, -1))] = "3"
+
+    starts = g.find("0") + g.find("1") + g.find("2") + g.find("3")
+    graph = defaultdict(set)
+    starts_seen = set()
+
+    while starts:
+        start = starts.pop()
+        if start in starts_seen:
+            continue
+        else:
+            starts_seen.add(start)
+
+        start_symbol = g[start]
+        xs = [start]
+        seen = set()
+
+        for steps in range(1_000):
+            nxs = []
+
+            for x in xs:
+                if x in seen:
+                    continue
+                else:
+                    seen.add(x)
+
+                for nb in g.neighbors_ort(x):
+                    if g[nb] == ".":
+                        nxs.append(nb)
+                    elif g[nb] in LETTERS:
+                        symbol = g[nb]
+                        if symbol != start_symbol:
+                            graph[start].add((nb, steps + 1))
+                            graph[nb].add((start, steps + 1))
+                        starts.append(nb)
+            xs = nxs
+            if len(xs) == 0:
+                break
+    # g.print()
+
+    all_keys = [g[p] for p in g.find(LETTERS_LOWER)]
+
+    robots = tuple(g.find("0") + g.find("1") + g.find("2") + g.find("3"))
+    poss = [(0, tuple(), robots)]
+    best: dict[tuple[tuple, tuple], int] = {(tuple(), tuple()): 0}
+
+    min_dist = 10**9
+    while poss:
+        current_distance, keys, robots = poss.pop()
+
+        if len(keys) == len(all_keys):
+            min_dist = min(min_dist, current_distance)
+            # print(min_dist)
+            continue
+
+        # print(current_distance, keys, robots)
+        for robot_i in range(len(robots)):
+            robot = robots[robot_i]
+            # robot_symbol = g[robot]
+            for next_pos, distance in graph[robot]:
+                next_symbol = g[next_pos]
+
+                if next_symbol in LETTERS_UPPER and not next_symbol.lower() in keys:
+                    continue
+
+                if next_symbol in LETTERS_LOWER:
+                    new_keys = set(keys)
+                    new_keys.add(next_symbol)
+                    new_keys = tuple(sorted(new_keys))
+                else:
+                    new_keys = keys
+
+                new_distance = current_distance + distance
+
+                new_robots = list(robots)
+                new_robots[robot_i] = next_pos
+                new_robots = tuple(new_robots)
+
+                key = (new_keys, new_robots)
+                if (key not in best) or (key in best and best[key] > new_distance):
+                    best[key] = new_distance
+                    poss.append((new_distance, new_keys, new_robots))
+
+        poss = sorted(poss, key=lambda xs: (xs[0], -len(xs[1])), reverse=True)
+        poss = poss[-10000:]
+    print(min_dist)
+
+
+def main():
+    data = get_data(__file__)
+    part_1(data)
+    part_2(data)
+
+
+if __name__ == "__main__":
+    main()