advent-of-code/2022/day19.py
Mikael CAPELLE ddebd26db2 Tmp
2022-12-19 18:46:16 +01:00

488 lines
14 KiB
Python

# -*- encoding: utf-8 -*-
import heapq
import math
import sys
import time
from collections import defaultdict
from typing import Literal, TypedDict
import numpy as np
from tqdm import tqdm
Reagent = Literal["ore", "clay", "obsidian", "geode"]
REAGENTS: tuple[Reagent] = (
"ore",
"clay",
"obsidian",
"geode",
)
IntOfReagent = dict[Reagent, int]
lines = sys.stdin.read().splitlines()
blueprints: list[dict[Reagent, IntOfReagent]] = [
{
"ore": {"ore": 4},
"clay": {"ore": 2},
"obsidian": {"ore": 3, "clay": 14},
"geode": {"ore": 2, "obsidian": 7},
},
{
"ore": {"ore": 2},
"clay": {"ore": 3},
"obsidian": {"ore": 3, "clay": 8},
"geode": {"ore": 3, "obsidian": 12},
},
]
class State:
robots: IntOfReagent
reagents: IntOfReagent
def __init__(
self,
robots: IntOfReagent | None = None,
reagents: IntOfReagent | None = None,
):
if robots is None:
assert reagents is None
self.reagents = {reagent: 0 for reagent in REAGENTS}
self.robots = {reagent: 0 for reagent in REAGENTS}
self.robots["ore"] = 1
else:
assert robots is not None and reagents is not None
self.robots = robots
self.reagents = reagents
def __eq__(self, other) -> bool:
return (
isinstance(other, State)
and self.robots == other.robots
and self.reagents == other.reagents
)
def __lt__(self, other) -> bool:
return isinstance(other, State) and tuple(
(self.robots[r], self.reagents[r]) for r in REAGENTS
) > tuple((other.robots[r], other.reagents[r]) for r in REAGENTS)
def __hash__(self) -> int:
return hash(tuple((self.robots[r], self.reagents[r]) for r in REAGENTS))
def __str__(self) -> str:
return "State({}, {})".format(
"/".join(str(self.robots[k]) for k in REAGENTS),
"/".join(str(self.reagents[k]) for k in REAGENTS),
)
def __repr__(self) -> str:
return str(self)
def dominates(lhs: State, rhs: State):
return all(
lhs.robots[r] >= rhs.robots[r] and lhs.reagents[r] >= rhs.reagents[r]
for r in REAGENTS
)
MAX_TIME = 24
blueprint = blueprints[1]
# parents: dict[State, tuple[State | None, int]] = {State(): (None, 0)}
# queue = [(0, State())]
# visited: set[State] = set()
# at_time: dict[int, list[State]] = defaultdict(lambda: [])
# while queue:
# time, state = heapq.heappop(queue)
# if state in visited:
# continue
# print(time, state)
# visited.add(state)
# at_time[time].append(state)
# if time > MAX_TIME:
# continue
# if len(queue) % 200 == 0:
# print(len(queue), len(visited), time)
# can_build_any: bool = False
# for reagent in REAGENTS:
# needed = blueprint[reagent]
# if any(state.robots[r] == 0 for r in needed):
# continue
# time_to_complete = max(
# max(
# math.ceil((needed[r] - state.reagents[r]) / state.robots[r])
# for r in needed
# ),
# 0,
# )
# # if time_to_complete != 0:
# # continue
# if time + time_to_complete + 1 > MAX_TIME:
# continue
# wait = time_to_complete + 1
# reagents = {
# r: state.reagents[r] + wait * state.robots[r] - needed.get(r, 0)
# for r in REAGENTS
# }
# robots = state.robots.copy()
# robots[reagent] += 1
# state_2 = State(reagents=reagents, robots=robots)
# if state_2 in visited:
# continue
# if any(dominates(state_v, state_2) for state_v in at_time[time + wait]):
# continue
# # print(time + wait)
# # if any(dominates(state_3, state_2) for state_3 in at_time[time + wait]):
# # print("?")
# # continue
# if state_2 not in parents or parents[state_2][1] > time + wait:
# parents[state_2] = (state, time + wait)
# heapq.heappush(queue, (time + wait, state_2))
# can_build_any = True
# at_time[time + wait].append(state_2)
# if not can_build_any:
# state_2 = State(
# reagents={
# r: state.reagents[r] + state.robots[r] * (MAX_TIME - time)
# for r in REAGENTS
# },
# robots=state.robots,
# )
# if state_2 in visited:
# continue
# if state_2 not in parents or parents[state_2][1] > time + wait:
# parents[state_2] = (state, MAX_TIME)
# heapq.heappush(queue, (MAX_TIME, state_2))
# print(len(visited))
# print(max(state.reagents["geode"] for state in visited))
# exit()
# while states:
# state = states.pop()
# processed.append(state)
# if state.time > MAX_TIME:
# continue
# if len(states) % 100 == 0:
# print(len(states), len(processed), min((s.time for s in states), default=1))
# can_build_any: bool = False
# for reagent in REAGENTS:
# needed = blueprint[reagent]
# if any(state.robots[r] == 0 for r in needed):
# continue
# time_to_complete = max(
# max(
# math.ceil((needed[r] - state.reagents[r]) / state.robots[r])
# for r in needed
# ),
# 0,
# )
# if state.time + time_to_complete + 1 > MAX_TIME:
# continue
# wait = time_to_complete + 1
# reagents = {
# r: state.reagents[r] + wait * state.robots[r] - needed.get(r, 0)
# for r in REAGENTS
# }
# robots = state.robots.copy()
# robots[reagent] += 1
# can_build_any = True
# state_2 = State(time=state.time + wait, reagents=reagents, robots=robots)
# # print(f"{state} -> {state_2}")
# states.add(state_2)
# if not any(dominates(s2, state_2) for s2 in states):
# states.add(state)
# # print(f"can build {reagent} in {time_to_complete}")
# if not can_build_any:
# states.add(
# State(
# time=MAX_TIME + 1,
# reagents={
# r: state.reagents[r] + state.robots[r] * (MAX_TIME - state.time)
# for r in REAGENTS
# },
# robots=state.robots,
# )
# )
# if len(states) % 1000 == 0:
# print("filtering")
# states = {
# s1
# for s1 in states
# if not any(dominates(s2, s1) for s2 in states if s2 is not s1)
# }
# # if len(states) > 4:
# # break
# # break
# print(len(processed))
# print(max(state.reagents["geode"] for state in processed))
# exit()
# for t in range(1, 25):
# states = set()
# for state in state_after_t[t - 1]:
# robots_that_can_be_built = [
# robot
# for robot in REAGENTS
# if all(
# state.reagents[reagent] >= blueprint[robot].get(reagent, 0)
# for reagent in REAGENTS
# )
# ]
# new_states = set()
# # new reagents
# reagents = {
# reagent: state.reagents[reagent] + state.robots[reagent]
# for reagent in REAGENTS
# }
# # if we can build anything, there is no point in waiting
# if len(robots_that_can_be_built) != len(REAGENTS):
# new_states.add(State(robots=state.robots, reagents=reagents))
# for robot in robots_that_can_be_built:
# robots = state.robots.copy()
# robots[robot] += 1
# reagents = {
# reagent: state.reagents[reagent]
# + state.robots[reagent]
# - blueprint[robot].get(reagent, 0)
# for reagent in REAGENTS
# }
# new_states.add(State(robots=robots, reagents=reagents))
# new_states = [
# s1
# for s1 in new_states
# if not any(s1 is not s2 and dominates(s2, s1) for s2 in new_states)
# ]
# states = {
# s1 for s1 in states if not any(dominates(s2, s1) for s2 in new_states)
# }
# states.update(new_states)
# state_after_t[t] = states
# exit()
MAX_TIME = 24
blueprint = blueprints[0]
state_after_t: dict[int, list[State]] = {0: [State()]}
for t in range(1, 25):
print(t, len(state_after_t[t - 1]))
bests_for_robots: dict[tuple[int, ...], list[State]] = {}
bests_for_reagents: dict[tuple[int, ...], list[State]] = {}
state_after_t[t] = []
t1 = time.time()
for state in state_after_t[t - 1]:
robots_that_can_be_built = [
robot
for robot in REAGENTS
if all(
state.reagents[reagent] >= blueprint[robot].get(reagent, 0)
for reagent in REAGENTS
)
]
# print(t, robots_that_can_be_built)
new_states: set[State] = set()
# new reagents
reagents = {
reagent: state.reagents[reagent] + state.robots[reagent]
for reagent in REAGENTS
}
# if we can build anything, there is no point in waiting
new_states.add(State(robots=state.robots, reagents=reagents))
for robot in robots_that_can_be_built:
robots = state.robots.copy()
robots[robot] += 1
reagents = {
reagent: state.reagents[reagent]
+ state.robots[reagent]
- blueprint[robot].get(reagent, 0)
for reagent in REAGENTS
}
new_states.add(State(robots=robots, reagents=reagents))
for s1 in new_states:
r1 = tuple(s1.robots[r] for r in REAGENTS)
if r1 not in bests_for_robots:
bests_for_robots[r1] = [s1]
else:
is_dominated = False
for s2 in bests_for_robots[r1]:
if all(s2.reagents[r] >= s1.reagents[r] for r in REAGENTS):
is_dominated = True
break
if not is_dominated:
bests_for_robots[r1].append(s1)
r2 = tuple(s1.reagents[r] for r in REAGENTS)
if r2 not in bests_for_reagents:
bests_for_reagents[r2] = [s1]
else:
is_dominated = False
for s2 in bests_for_reagents[r2]:
if all(s2.robots[r] >= s1.robots[r] for r in REAGENTS):
is_dominated = True
break
if not is_dominated:
bests_for_reagents[r2].append(s1)
# state_after_t[t].extend(new_states)
t2 = time.time()
for bests in bests_for_robots.values():
dominated = [False for _ in range(len(bests))]
for i_s1, s1 in enumerate(bests):
if dominated[i_s1]:
continue
for i_s2, s2 in enumerate(bests[i_s1 + 1 :], start=i_s1 + 1):
if dominated[i_s2]:
continue
if all(s1.reagents[r] >= s2.reagents[r] for r in REAGENTS):
dominated[i_s2] = True
state_after_t[t].extend(
s1 for i_s1, s1 in enumerate(bests) if not dominated[i_s1]
)
for bests in bests_for_reagents.values():
dominated = [False for _ in range(len(bests))]
for i_s1, s1 in enumerate(bests):
if dominated[i_s1]:
continue
for i_s2, s2 in enumerate(bests[i_s1 + 1 :], start=i_s1 + 1):
if dominated[i_s2]:
continue
if all(s1.robots[r] >= s2.robots[r] for r in REAGENTS):
dominated[i_s2] = True
state_after_t[t].extend(
s1 for i_s1, s1 in enumerate(bests) if not dominated[i_s1]
)
t3 = time.time()
np_states = np.array(
[
[state.robots[r] for r in REAGENTS] + [state.reagents[r] for r in REAGENTS]
for state in state_after_t[t]
]
)
dominated = np.zeros(len(np_states), dtype=bool)
t4 = time.time()
# c = (np_states[None, :, :] <= np_states[:, None, :]).all(axis=-1)
# c[np.arange(len(np_states)), np.arange(len(np_states))] = False
# dominated = c.any(axis=0)
for i in range(len(np_states)):
if dominated[i]:
continue
dominated[i] = not (np_states[i + 1 :] <= np_states[i]).any(axis=1)
dominated[i + 1 :] = (np_states[i + 1 :] <= np_states[i]).all(axis=1)
t5 = time.time()
state_after_t[t] = list(np.array(state_after_t[t])[~dominated])
t6 = time.time()
print(
"->",
t,
len(state_after_t[t]),
dominated.sum(),
t2 - t1,
t3 - t2,
t4 - t3,
t5 - t4,
t6 - t5,
)
# print("->", len(state_after_t[t]))
# dominated = [False for _ in range(len(state_after_t[t]))]
# keep = set()
# for i_s1, s1 in enumerate(tqdm(state_after_t[t])):
# if dominated[i_s1]:
# continue
# for i_s2, s2 in enumerate(state_after_t[t][i_s1 + 1 :], start=i_s1 + 1):
# if dominated[i_s2]:
# continue
# if dominates(s1, s2):
# dominated[i_s2] = True
# elif dominates(s2, s1):
# dominated[i_s1] = True
# break
# if not dominated[i_s1]:
# keep.add(s1)
# state_after_t[t] = list(keep)
# print(len(state_after_t[t]))
# print(sum(dominated))
# break
print(max(state.reagents["geode"] for state in state_after_t[24]))