134 lines
3.3 KiB
Python
134 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from functools import total_ordering
|
|
from collections import defaultdict
|
|
from itertools import combinations
|
|
|
|
@total_ordering
|
|
class MaxBound:
|
|
def __gt__(self, other):
|
|
return True
|
|
|
|
def __eq__(self, other):
|
|
return False
|
|
|
|
def __ge__(self, other):
|
|
return True
|
|
|
|
def __le__(self, other):
|
|
return False
|
|
|
|
def __repr__(self):
|
|
return "MaxBound"
|
|
|
|
|
|
def move_func(maze, i):
|
|
w = maze.index('\n')+1
|
|
possible = i+1, i-1, i+w, i-w
|
|
return [j for j in possible if j < len(maze) and maze[j] not in ('\n', '#')]
|
|
|
|
def build_branch(maze, i, w=0, visited=None, edges=None):
|
|
if visited == None:
|
|
visited = defaultdict(MaxBound)
|
|
edges = defaultdict(MaxBound)
|
|
if visited[i] <= w:
|
|
return
|
|
visited[i] = w
|
|
node = maze[i]
|
|
if node != '.' and w > 0:
|
|
edges[node] = min(edges[node], w)
|
|
return
|
|
for m in move_func(maze, i):
|
|
build_branch(maze, m, w+1, visited, edges)
|
|
return edges
|
|
|
|
def build_graph(maze):
|
|
keys = [c for c in maze if c not in ('#', '.', '\n')]
|
|
graph = dict()
|
|
for key in keys:
|
|
graph[key] = build_branch(maze, maze.index(key))
|
|
return graph
|
|
|
|
def print_graph(graph):
|
|
for node in graph:
|
|
print(node, end='')
|
|
for n, w in graph[node].items():
|
|
print("--{}--> {}".format(w, n))
|
|
print(' '*len(node), end='')
|
|
print()
|
|
|
|
def unlock(graph, key):
|
|
new = dict()
|
|
if key not in graph:
|
|
return graph
|
|
for a in graph:
|
|
if a == key:
|
|
continue
|
|
new[a] = defaultdict(MaxBound)
|
|
for b, w in graph[a].items():
|
|
if b == key:
|
|
continue
|
|
new[a][b] = w
|
|
for a, b in combinations(graph[key].keys(), 2):
|
|
w = graph[key][a] + graph[key][b]
|
|
new[a][b] = new[b][a] = min(new[a][b], w)
|
|
return new
|
|
|
|
def is_key(c):
|
|
return c == c.lower()
|
|
|
|
def doorof(key):
|
|
return key.upper() if key.isalpha() else None
|
|
|
|
def heu_sort(edges):
|
|
return sorted(edges.items(), key=lambda x: x[1])
|
|
|
|
def search(graph, nodes, w=0, best=MaxBound(), d=0, mem=None, key=''):
|
|
if mem == None:
|
|
mem = defaultdict(MaxBound)
|
|
|
|
id = "".join(sorted(nodes)) + ''.join(sorted(n for n in graph if n not in nodes and is_key(n)))
|
|
if w >= best or w >= mem[id]:
|
|
return best
|
|
|
|
if len(mem) < 10**5:
|
|
mem[id] = w
|
|
|
|
graph = unlock(graph, doorof(key))
|
|
|
|
for i, node in enumerate(nodes):
|
|
bridged = unlock(graph, node)
|
|
|
|
for n, w2 in heu_sort(graph[node]):
|
|
if not is_key(n):
|
|
continue
|
|
r = search(bridged, nodes[:i] + [n] + nodes[i+1:], w+w2, best, d+1, mem, n)
|
|
best = min(best, r)
|
|
if len(graph) == len(nodes) and w < best:
|
|
best = w
|
|
return best
|
|
|
|
def init_maze(maze):
|
|
maze = list(maze)
|
|
for i in range(maze.count('@')):
|
|
maze[maze.index('@')] = str(i)
|
|
return ''.join(maze), map(str, range(i+1))
|
|
|
|
def prepoc(puzzle_input):
|
|
return puzzle_input
|
|
|
|
def partI(maze):
|
|
graph = build_graph(maze)
|
|
return search(graph, ['@'])
|
|
|
|
def partII(maze):
|
|
maze = list(maze)
|
|
i = maze.index('@')
|
|
w = maze.index('\n')
|
|
maze[i-1:i+2] = '###'
|
|
maze[i-w-2:i-w+1] = maze[i+w:i+w+3] = '@#@'
|
|
maze, nodes = init_maze(maze)
|
|
graph = build_graph(maze)
|
|
return search(graph, list(nodes))
|
|
|