Files
aocpy/2024/d21.py

264 lines
7.4 KiB
Python

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)