diff --git a/src/holt59/aoc/2022/day16.py b/src/holt59/aoc/2022/day16.py index f9068a5..3b54ce2 100644 --- a/src/holt59/aoc/2022/day16.py +++ b/src/holt59/aoc/2022/day16.py @@ -6,8 +6,6 @@ import re from collections import defaultdict from typing import Any, FrozenSet, Iterator, NamedTuple -from tqdm import tqdm - from ..base import BaseSolver @@ -62,68 +60,70 @@ def update_with_better( node_at_times[flowing] = max(node_at_times[flowing], flow) -def part_1( - start_pipe: Pipe, - max_time: int, - distances: dict[tuple[Pipe, Pipe], int], - relevant_pipes: FrozenSet[Pipe], -): - node_at_times: dict[int, dict[Pipe, dict[FrozenSet[Pipe], int]]] = defaultdict( - lambda: defaultdict(lambda: defaultdict(lambda: 0)) - ) - node_at_times[0] = {start_pipe: {frozenset(): 0}} - - for time in range(max_time): - for c_pipe, nodes in node_at_times[time].items(): - for flowing, flow in nodes.items(): - for target in relevant_pipes: - distance = distances[c_pipe, target] + 1 - if time + distance >= max_time or target in flowing: - continue - - update_with_better( - node_at_times[time + distance][target], - flow + sum(pipe.flow for pipe in flowing) * distance, - flowing | {target}, - ) - - update_with_better( - node_at_times[max_time][c_pipe], - flow + sum(pipe.flow for pipe in flowing) * (max_time - time), - flowing, - ) - - return max( - flow - for nodes_of_pipe in node_at_times[max_time].values() - for flow in nodes_of_pipe.values() - ) - - -def part_2( - start_pipe: Pipe, - max_time: int, - distances: dict[tuple[Pipe, Pipe], int], - relevant_pipes: FrozenSet[Pipe], -): - def compute(pipes_for_me: FrozenSet[Pipe]) -> int: - return part_1(start_pipe, max_time, distances, pipes_for_me) + part_1( - start_pipe, max_time, distances, relevant_pipes - pipes_for_me - ) - - combs = [ - frozenset(relevant_pipes_1) - for r in range(2, len(relevant_pipes) // 2 + 1) - for relevant_pipes_1 in itertools.combinations(relevant_pipes, r) - ] - - return max(compute(comb) for comb in tqdm(combs)) - - # === MAIN === class Solver(BaseSolver): + def part_1( + self, + start_pipe: Pipe, + max_time: int, + distances: dict[tuple[Pipe, Pipe], int], + relevant_pipes: FrozenSet[Pipe], + ): + node_at_times: dict[int, dict[Pipe, dict[FrozenSet[Pipe], int]]] = defaultdict( + lambda: defaultdict(lambda: defaultdict(lambda: 0)) + ) + node_at_times[0] = {start_pipe: {frozenset(): 0}} + + for time in range(max_time): + for c_pipe, nodes in node_at_times[time].items(): + for flowing, flow in nodes.items(): + for target in relevant_pipes: + distance = distances[c_pipe, target] + 1 + if time + distance >= max_time or target in flowing: + continue + + update_with_better( + node_at_times[time + distance][target], + flow + sum(pipe.flow for pipe in flowing) * distance, + flowing | {target}, + ) + + update_with_better( + node_at_times[max_time][c_pipe], + flow + sum(pipe.flow for pipe in flowing) * (max_time - time), + flowing, + ) + + return max( + flow + for nodes_of_pipe in node_at_times[max_time].values() + for flow in nodes_of_pipe.values() + ) + + def part_2( + self, + start_pipe: Pipe, + max_time: int, + distances: dict[tuple[Pipe, Pipe], int], + relevant_pipes: FrozenSet[Pipe], + ): + def compute(pipes_for_me: FrozenSet[Pipe]) -> int: + return self.part_1( + start_pipe, max_time, distances, pipes_for_me + ) + self.part_1( + start_pipe, max_time, distances, relevant_pipes - pipes_for_me + ) + + combs = [ + frozenset(relevant_pipes_1) + for r in range(2, len(relevant_pipes) // 2 + 1) + for relevant_pipes_1 in itertools.combinations(relevant_pipes, r) + ] + + return max(compute(comb) for comb in self.progress.wrap(combs)) + def solve(self, input: str) -> Iterator[Any]: lines = [line.strip() for line in input.splitlines()] @@ -153,7 +153,7 @@ class Solver(BaseSolver): relevant_pipes = frozenset(pipe for pipe in pipes.values() if pipe.flow > 0) # 1651, 1653 - yield part_1(pipes["AA"], 30, distances, relevant_pipes) + yield self.part_1(pipes["AA"], 30, distances, relevant_pipes) # 1707, 2223 - yield part_2(pipes["AA"], 26, distances, relevant_pipes) + yield self.part_2(pipes["AA"], 26, distances, relevant_pipes)