159 lines
4.9 KiB
Python
159 lines
4.9 KiB
Python
import sys
|
|
import matrix
|
|
import shared
|
|
import scanf
|
|
from dataclasses import dataclass
|
|
from collections import defaultdict
|
|
from typing import List, Dict
|
|
from pprint import pprint
|
|
import networkx as nx
|
|
|
|
from IPython.display import Image, display
|
|
|
|
@dataclass
|
|
class Valve:
|
|
label: str
|
|
rate: int
|
|
tunnels: List[str]
|
|
opened_at: int = -1
|
|
potential: Dict[str, int] = None
|
|
|
|
def set_potential(self, valves):
|
|
self.potential = {}
|
|
for tunnel in self.tunnels:
|
|
self.potential[tunnel] = valves[tunnel].rate
|
|
|
|
def highest_potential(self):
|
|
return max(self.potential, key=self.potential.get)
|
|
|
|
|
|
def parse(rows):
|
|
valves = {}
|
|
for row in rows:
|
|
left, right = row.split(" valve")
|
|
right = right.replace("s ", "").lstrip()
|
|
valve, rate = scanf.scanf("Valve %s has flow rate=%d; %*s %*s to", left)
|
|
tunnels = right.split(", ")
|
|
valves[valve] = Valve(label=valve, rate=rate, tunnels=tunnels)
|
|
|
|
for _, v in valves.items():
|
|
v.set_potential(valves)
|
|
return valves
|
|
|
|
|
|
def part1(rows, sample=False):
|
|
p1 = Part1(rows, sample, 30)
|
|
p1.run()
|
|
|
|
|
|
class Part1:
|
|
def __init__(self, rows, sample, minutes):
|
|
self.rows = rows
|
|
self.sample = sample
|
|
self.valves = parse(rows)
|
|
self.nonzero = {v.label: v for _, v in self.valves.items() if v.rate > 0}
|
|
self.cur = self.valves["AA"]
|
|
self.tick = 1
|
|
self.minutes = minutes
|
|
self.g = nx.DiGraph()
|
|
self.path_distances = defaultdict(dict)
|
|
self.set_up_graph()
|
|
|
|
def draw(self):
|
|
pdot = nx.drawing.nx_pydot.to_pydot(self.g)
|
|
pdot.write_png("15.png")
|
|
|
|
def set_up_graph(self):
|
|
for lbl, v in self.valves.items():
|
|
for t in v.tunnels:
|
|
# self.g.add_edge(lbl, t, {'weight':self.valves[t].rate})
|
|
self.g.add_edge(lbl, t, weight=self.valves[t].rate)
|
|
all_keys = self.valves.keys()
|
|
for lbl, _ in self.valves.items():
|
|
for other in all_keys:
|
|
if other == lbl:
|
|
continue
|
|
self.path_distances[lbl][other] = min([len(x) for x in nx.shortest_simple_paths(self.g, lbl, other)])
|
|
|
|
def do_tick(self, minute):
|
|
pressure = 0
|
|
opened = []
|
|
for _, valve in self.valves.items():
|
|
if valve.opened_at > 0:
|
|
pressure += valve.rate
|
|
opened.append(valve.label)
|
|
print(f"== Min {minute}:: Valves {', '.join(opened)} are open, releasing {pressure} pressure")
|
|
|
|
def calculate_total_flow(self):
|
|
total = 0
|
|
for label, valve in self.valves.items():
|
|
if valve.opened_at > 0:
|
|
total += valve.rate * (30 - valve.opened_at)
|
|
return total
|
|
|
|
def run(self):
|
|
# Construct the graph with vertices & edges from the input
|
|
# Call a function to compute the distances between every pair of vertices
|
|
# Create a closed set containing all the valves with non-zero rates
|
|
# At each step, iterate over the remaining set of closed, non-zero valves
|
|
# - Subtract the distance from remaining minutes
|
|
# - Calculate the flow (rate * remaining minutes)
|
|
# - Remove the recently opened valve from the closed set (functionally), so the deeper levels won't consider it
|
|
|
|
def priority(remaining):
|
|
_pris = []
|
|
for _,n in self.nonzero.items():
|
|
# (time_remaining - distance_to_valve - 1) * flow rate
|
|
pri = (remaining - self.path_distances[self.cur.label][n.label] - 1) * n.rate
|
|
_pris.append((n.label, pri))
|
|
_pris = list(sorted(_pris, key=lambda x: x[1], reverse=True))
|
|
return _pris
|
|
|
|
|
|
remaining = self.minutes
|
|
open_order = []
|
|
while len(self.nonzero):
|
|
if remaining <= 0:
|
|
print("ran out of time")
|
|
break
|
|
self.do_tick(31-remaining)
|
|
pris = priority(remaining)
|
|
print(pris)
|
|
_pri, _ = pris.pop(0)
|
|
n = self.nonzero[_pri]
|
|
del self.nonzero[_pri]
|
|
open_order.append(n.label)
|
|
print("\tMoving to", n.label)
|
|
print("\tOpening ", n.label)
|
|
distance = self.path_distances[self.cur.label][n.label]
|
|
self.cur = n
|
|
self.cur.opened_at = self.minutes - (remaining + 1)
|
|
remaining -= distance # Move
|
|
remaining -= 1 # open
|
|
|
|
|
|
print(remaining)
|
|
print(open_order)
|
|
print("sample: 1651")
|
|
print("total flow:", self.calculate_total_flow())
|
|
|
|
|
|
|
|
def main():
|
|
sample = False
|
|
if sys.argv[-1] == "--sample":
|
|
sample = True
|
|
rows = [row for row in shared.load_rows(16)]
|
|
|
|
with shared.elapsed_timer() as elapsed:
|
|
part1(rows, sample)
|
|
print("🕒", elapsed())
|
|
|
|
# with shared.elapsed_timer() as elapsed:
|
|
# part2(rows, sample)
|
|
# print("🕒", elapsed())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|