import heapq from lib import get_data data = get_data(__file__) data = """029A 980A 179A 456A 379A""" 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 def find_numpad(target: str) -> str: COUNT, CURRENT, RESULT, PATH = 0, 1, 2, 3 def dist(n1, n2): if n1 == 0: return 0 return 1 def h(node): return len(target) - len(node[RESULT]) def is_goal(node): return node[RESULT] == target def neighbors(node): nbs = [] for press in 'A^>v<': nnode = list(node) nnode[PATH] = node[PATH] + press if press == "A": nnode[RESULT] = node[RESULT] + node[CURRENT] else: nnode[CURRENT] = new_num(node[CURRENT], press) if nnode[CURRENT] is None: continue if not target.startswith(nnode[RESULT]): continue nnode[COUNT] += 1 nbs.append(tuple(nnode)) return nbs start = (0, 'A', '', '') starts = [start] open_set = [] g_score = {} cost = None solutions = [] 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] cost = current[PATH] if solutions != [] and len(cost) > len(solutions[-1]): break solutions.append(current[PATH]) 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 solutions def find_dirpad(target: str) -> str: COUNT, CURRENT, RESULT, PATH = 0, 1, 2, 3 def dist(n1, n2): if n1 == 0: return 0 return 1 def h(node): hdist = { "A": {"^": 2, ">": 2, "<": 4, "v": 3, "A": 1}, ">": {"^": 3, ">": 1, "<": 3, "v": 2, "A": 2}, "v": {"^": 2, ">": 2, "<": 2, "v": 1, "A": 3}, "<": {"^": 3, ">": 3, "<": 1, "v": 2, "A": 4}, "^": {"^": 1, ">": 3, "<": 3, "v": 2, "A": 2}, } d = 0 cur_key = node[CURRENT] nxt_key_idx = len(node[RESULT]) while nxt_key_idx < len(target): nxt_key = target[nxt_key_idx] d += hdist[cur_key][nxt_key] cur_key = nxt_key nxt_key_idx += 1 return d def is_goal(node): return node[RESULT] == target def neighbors(node): nbs = [] to_press = list('A^>v<') if node[PATH]: to_press.remove(node[PATH][-1]) to_press = [node[PATH][-1]] + to_press for press in to_press: if press == "A": new_result = node[RESULT] + node[CURRENT] new_current = node[CURRENT] if not target.startswith(new_result): continue else: new_result = node[RESULT] new_current = new_dir(node[CURRENT], press) if new_current is None: continue nnode = (node[COUNT] + 1, new_current, new_result, node[PATH] + press) nbs.append(nnode) return nbs start = (0, 'A', '', '') starts = [start] open_set = [] g_score = {} cost = None solutions = [] 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] cost = current[PATH] if solutions and len(cost) > len(solutions[-1]): break solutions.append(cost) 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 solutions def find(a): print(a) a = find_numpad(a)[0] print(len(a), a) a = find_dirpad(a)[0] print(len(a), a) a = find_dirpad(a)[0] print(len(a), a) return a print(a) for a in find_numpad(a): print(len(a), a) for a in find_dirpad(a): print(len(a), a) a = find_dirpad(a) # print(len(a)) # print() exit() # a = find_keypad(a) return a t = 0 for line in data.splitlines(): s = find(line) i = int("".join([c for c in line if ord('0') <= ord(c) <= ord('9')])) print(len(s), i, s) # t += s * i print(t)