Fix 2022 day 16 for progress API.

This commit is contained in:
Mikaël Capelle 2024-12-08 19:25:21 +01:00
parent 9d7ef94fa6
commit dd3f332870

View File

@ -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)