104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
|
#!/bin/python3
|
||
|
|
|
||
|
|
import re
|
||
|
|
from dijkstar import Graph, find_path
|
||
|
|
import numpy as np
|
||
|
|
from itertools import combinations
|
||
|
|
import itertools as it
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
valves = list()
|
||
|
|
|
||
|
|
with open('inputs/day16.txt', 'r') as f:
|
||
|
|
for line in f:
|
||
|
|
m = re.match(r"Valve (?P<name>[A-Z]{2}) has flow rate=(?P<rate>\d+); tunnels? leads? to valves? (?P<tunnels>.*)", line.strip())
|
||
|
|
name = m.group("name")
|
||
|
|
rate = int(m.group("rate"))
|
||
|
|
tunnels = m.group("tunnels").split(", ")
|
||
|
|
|
||
|
|
valves.append({"name": name, "rate": rate, "tunnels": tunnels})
|
||
|
|
|
||
|
|
indices = {}
|
||
|
|
for i, valve in enumerate(valves):
|
||
|
|
indices[valve["name"]] = i
|
||
|
|
|
||
|
|
graph = Graph()
|
||
|
|
for v1 in valves:
|
||
|
|
for v2 in v1["tunnels"]:
|
||
|
|
graph.add_edge(v1["name"], v2, 1)
|
||
|
|
graph.add_edge(v2, v1["name"], 1)
|
||
|
|
|
||
|
|
adjs = np.zeros((len(valves), len(valves)))
|
||
|
|
for v1, v2 in combinations(map(lambda v: v["name"], valves), 2):
|
||
|
|
cost = find_path(graph, v1, v2).total_cost
|
||
|
|
adjs[indices[v1], indices[v2]] = cost
|
||
|
|
adjs[indices[v2], indices[v1]] = cost
|
||
|
|
|
||
|
|
class Valve:
|
||
|
|
def __init__(self, name, rate):
|
||
|
|
self.name = name
|
||
|
|
self.rate = rate
|
||
|
|
|
||
|
|
def value(self, start, time_left):
|
||
|
|
cost = adjs[indices[start], indices[self.name]] + 1
|
||
|
|
return (self.rate * (time_left - cost), cost) # one off?
|
||
|
|
|
||
|
|
valve_map = dict()
|
||
|
|
valve_set = set()
|
||
|
|
for valve in valves:
|
||
|
|
if valve["rate"] > 0:
|
||
|
|
valve_set.add(valve["name"])
|
||
|
|
valve_map[valve["name"]] = Valve(valve["name"], valve["rate"])
|
||
|
|
|
||
|
|
def search(valve_set, current, time_left, visited):
|
||
|
|
to_visit_set = valve_set.difference(visited)
|
||
|
|
best_score = 0
|
||
|
|
for to_visit in to_visit_set:
|
||
|
|
# print("time left:", time_left)
|
||
|
|
# print("current:", current)
|
||
|
|
# print("going to visit:", to_visit)
|
||
|
|
visited_copy = visited.copy()
|
||
|
|
visited_copy.add(to_visit)
|
||
|
|
score, dtime = valve_map[to_visit].value(current, time_left)
|
||
|
|
# print("+score:", score, "+time", dtime)
|
||
|
|
if time_left - dtime >= 0:
|
||
|
|
score = score + search(valve_set, to_visit, time_left - dtime, visited_copy)
|
||
|
|
if not best_score or score > best_score:
|
||
|
|
best_score = score
|
||
|
|
return best_score
|
||
|
|
|
||
|
|
|
||
|
|
# print('part1', search(valve_set, "AA", 30, set()))
|
||
|
|
|
||
|
|
# part2 idea
|
||
|
|
# BIGBRAIN
|
||
|
|
# divide valve set in two disjoint sets
|
||
|
|
# apply same search algorithm to both sets.
|
||
|
|
# sum the scores of the sets to get total score
|
||
|
|
# brute force set distribution
|
||
|
|
|
||
|
|
def partition(pred, iterable):
|
||
|
|
t1, t2 = it.tee(iterable)
|
||
|
|
return it.filterfalse(pred, t1), filter(pred, t2)
|
||
|
|
|
||
|
|
def partitions(l):
|
||
|
|
return list(filter(lambda x: [] not in x, [[[x[1] for x in f] for f in partition(lambda x: x[0], zip(pattern, l))] for pattern in it.product([True, False], repeat=len(l))]))
|
||
|
|
# for [lhs, rhs] in parts:
|
||
|
|
# if (set(lhs), set(rhs)) not in result and (set(rhs), set(lhs)) not in result:
|
||
|
|
# result.append((set(lhs), set(rhs)))
|
||
|
|
# return result
|
||
|
|
|
||
|
|
best_score = 0
|
||
|
|
|
||
|
|
for part in tqdm(partitions(list(valve_set))):
|
||
|
|
[part1, part2] = part
|
||
|
|
part1 = set(part1)
|
||
|
|
score1 = search(part1, 'AA', 26, set())
|
||
|
|
part2 = set(part2)
|
||
|
|
score2 = search(part2, 'AA', 26, set())
|
||
|
|
grand_score = score1 + score2
|
||
|
|
if grand_score > best_score:
|
||
|
|
best_score = grand_score
|
||
|
|
|
||
|
|
print('part2', best_score)
|