aoc/22/day16.py

104 lines
3.2 KiB
Python
Raw Normal View History

2023-04-17 18:18:20 +00:00
#!/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)