2024 d21 part 1

2024/d21.py

@@ -0,0 +1,189 @@
+import heapq
+from lib import get_data
+from functools import cache
+
+
+data = get_data(__file__)
+
+
+def new_dir(current, move):
+    if current == "<":
+        if move == ">": return "v"
+        else: return None
+    elif current == "^":
+        if move == ">": return "A"
+        elif move == "v": return "v"
+        else: return None
+    elif current == "v":
+        if move == "<": return "<"
+        elif move == "^": return "^"
+        elif move == ">": return ">"
+        else: return None
+    elif current == ">":
+        if move == "<": return "v"
+        elif move == "^": return "A"
+        else: return None
+    elif current == "A":
+        if move == "<": return "^"
+        elif move == "v": return ">"
+        else: return None
+    else:
+        assert False
+
+def new_num(current, move):
+    if current == "7":
+        if move == ">": return "8"
+        elif move == "v": return "4"
+        else: return None
+    elif current == "8":
+        if move == "<": return "7"
+        elif move == ">": return "9"
+        elif move == "v": return "5"
+        else: return None
+    elif current == "9":
+        if move == "<": return "8"
+        elif move == "v": return "6"
+        else: return None
+    elif current == "5":
+        if move == "^": return "8"
+        elif move == "<": return "4"
+        elif move == ">": return "6"
+        elif move == "v": return "2"
+        else: assert False
+    elif current == "4":
+        if move == "^": return "7"
+        elif move == ">": return "5"
+        elif move == "v": return "1"
+        else: return None
+    elif current == "6":
+        if move == "^": return "9"
+        elif move == "<": return "5"
+        elif move == "v": return "3"
+        else: return None
+    elif current == "1":
+        if move == "^": return "4"
+        elif move == ">": return "2"
+        else: return None
+    elif current == "2":
+        if move == "^": return "5"
+        elif move == "<": return "1"
+        elif move == ">": return "3"
+        elif move == "v": return "0"
+        else: assert False
+    elif current == "3":
+        if move == "^": return "6"
+        elif move == "<": return "2"
+        elif move == "v": return "A"
+        else: return None
+    elif current == "0":
+        if move == "^": return "2"
+        elif move == ">": return "A"
+        else: return None
+    elif current == "A":
+        if move == "^": return "3"
+        elif move == "<": return "0"
+        else: return None
+    else:
+        assert False
+
+
+@cache
+def new_state(press, node):
+    cost, r, n, path = node
+    r = list(r)
+
+    for i in range(len(r)):
+        if press == "A":
+            press = r[i]
+        else:
+            r[i] = new_dir(r[i], press)
+            if r[i] is None:
+                return None
+            press = None
+            break
+
+    npath = str(path)
+    if press is not None:
+        if press == "A":
+            npath += n
+        else:
+            n = new_num(n, press)
+            if n is None:
+                return None
+
+    nnode = cost + 1, tuple(r), n, npath
+    return nnode
+
+    if press == "A":
+        if r[0] == "A":
+            if r[1] == "A":
+                path += n
+            else:
+                n = new_num(n, r[1])
+        else:
+            r[1] = new_dir(r[1], r[0])
+    else:
+        r[0] = new_dir(r[0], press)
+
+    if r[0] is None or r[1] is None or n is None:
+        return None
+
+    npath = str(path)
+    return cost + 1, tuple(r), n, npath
+
+def sequence(line):
+    def dist(n1, n2):
+        """cost from node to node"""
+        return 1
+
+    def h(node):
+        """heuristic function (never overestimate)"""
+        return len(line) - len(node[3])
+
+    def is_goal(node):
+        return node[3] == line
+
+    def neighbors(node):
+        nbs = []
+        for move in '<^v>A':
+            nb = new_state(move, node)
+            if nb is None:
+                continue
+            if not line.startswith(node[3]):
+                continue
+            nbs.append(nb)
+        return nbs
+
+    key_pads = tuple(['A' for _ in range(5)])
+    start = (0, key_pads, 'A', '')
+    starts = [start]
+    open_set = []
+    g_score = {}
+    cost = None
+
+    for start in starts:
+        heapq.heappush(open_set, (h(start), start))
+        g_score[start] = dist(0, start)
+
+    while open_set:
+        current_f_score, current = heapq.heappop(open_set)
+        if is_goal(current):
+            assert current_f_score == g_score[current]
+            cost = g_score[current]
+            break
+
+        for neighbor in neighbors(current):
+            tentative_g_score = g_score[current] + dist(current, neighbor)
+            if neighbor not in g_score or tentative_g_score < g_score[neighbor]:
+                g_score[neighbor] = tentative_g_score
+                f_score = g_score[neighbor] + h(neighbor)
+                heapq.heappush(open_set, (f_score, neighbor))
+    return cost
+
+t = 0
+for line in data.splitlines():
+    s = sequence(line)
+    i = int("".join([c for c in line if ord('0') <= ord(c) <= ord('9')]))
+    print(s, i)
+    t += (s - 1) * i
+print(t)