115 lines
2.6 KiB
Python
115 lines
2.6 KiB
Python
import re
|
|
import sys
|
|
|
|
# PART 1
|
|
|
|
STEP_RE = r"Step (\S+) must be finished before step (\S+) can begin.\n"
|
|
|
|
def load_steps(filename):
|
|
steps = {}
|
|
with open(filename, "r") as f:
|
|
for line in f:
|
|
match = re.fullmatch(STEP_RE, line)
|
|
step, before = match.groups()
|
|
steps[step] = steps.get(step, set()) | {before}
|
|
|
|
return steps
|
|
|
|
def reverse_steps(steps):
|
|
reverse = {}
|
|
for step, befores in steps.items():
|
|
# Make sure that step exists in reverse
|
|
reverse[step] = reverse.get(step, set())
|
|
|
|
for before in befores:
|
|
reverse[before] = reverse.get(before, set()) | {step}
|
|
|
|
return reverse
|
|
|
|
def duration_of(step):
|
|
return ord(step) - ord("A") + 61
|
|
|
|
class Tree:
|
|
def __init__(self, steps, workers=1):
|
|
self.workers = {i: None for i in range(workers)}
|
|
self.steps = reverse_steps(steps) # Warning: Steps are reversed in Trees.
|
|
self.result = []
|
|
self.duration = 0
|
|
|
|
def working(self):
|
|
return {worker: work for worker, work in self.workers.items() if work is not None}
|
|
|
|
def find_free(self):
|
|
return {step for step, afters in self.steps.items() if len(afters) == 0}
|
|
|
|
def find_working(self):
|
|
return {step for (step, _) in self.working().values()}
|
|
|
|
def find_available(self):
|
|
return self.find_free() - self.find_working()
|
|
|
|
def remove_step(self, step):
|
|
try:
|
|
del self.steps[step]
|
|
except KeyError:
|
|
pass
|
|
|
|
for s in self.steps.values():
|
|
try:
|
|
s.remove(step)
|
|
except KeyError:
|
|
pass
|
|
|
|
def update_workers(self):
|
|
min_duration = min(duration for (_, duration) in self.working().values())
|
|
self.duration += min_duration
|
|
|
|
finished_steps = set()
|
|
|
|
# Subtract min_duration from all workers
|
|
for w, s in self.workers.items():
|
|
if s is not None:
|
|
step, duration = s
|
|
duration -= min_duration
|
|
|
|
if duration <= 0:
|
|
finished_steps.add(step)
|
|
self.remove_step(step)
|
|
self.workers[w] = None
|
|
else:
|
|
self.workers[w] = (step, duration)
|
|
|
|
self.result += list(finished_steps)
|
|
|
|
def distribute_jobs(self):
|
|
available = list(reversed(sorted(self.find_available())))
|
|
|
|
for w, s in self.workers.items():
|
|
if not available:
|
|
break
|
|
|
|
if s is None:
|
|
step = available.pop()
|
|
duration = duration_of(step)
|
|
self.workers[w] = (step, duration)
|
|
|
|
def run(self):
|
|
while self.steps:
|
|
self.distribute_jobs()
|
|
self.update_workers()
|
|
|
|
def main(filename):
|
|
print(f"Solutions for {filename}")
|
|
steps = load_steps(filename)
|
|
tree = Tree(steps, workers=1)
|
|
tree.run()
|
|
sequence = "".join(tree.result)
|
|
print(f"Part 1: {sequence}")
|
|
tree = Tree(steps, workers=5)
|
|
tree.run()
|
|
duration = tree.duration
|
|
print(f"Part 2: {duration}")
|
|
|
|
if __name__ == "__main__":
|
|
for filename in sys.argv[1:]:
|
|
main(filename)
|