aoc2022/d19/solve.py

33 lines
1.3 KiB
Python
Raw Normal View History

2022-12-19 23:02:42 +00:00
import sys, re, numpy
V = lambda *a: numpy.array(a)
k = lambda a: tuple(sum(a))
def parse(line):
i,a,b,c,d,e,f = map(int, re.findall(r'\d+',line))
return (i, (V(0,0,0,a), V(0,0,0,1)), # Cost and production
(V(0,0,0,b), V(0,0,1,0)), # of each robot type,
(V(0,0,d,c), V(0,1,0,0)), # in the order geode,
(V(0,f,0,e), V(1,0,0,0)), # obs, clay, and ore.
(V(0,0,0,0), V(0,0,0,0))) # Construct no robot.
BEST = int(sys.argv[1])
def run(blueprint, t):
todo = [(V(0,0,0,0), V(0,0,0,1))] # resources and robots.
for _ in range(t):
todo_ = list() # Queue for the next minute.
for have, make in todo:
for cost, more in blueprint:
if all(have >= cost): # We can afford this robot.
todo_.append((have+make-cost, make+more))
todo = sorted(todo_, key=k)[-BEST:] # adds resources and robots (k)
# the more of both, the better
# then prune the search queue
return max(todo, key=k)[0][0]
part1, part2 = 0, 1
for i, *blueprint in map(parse, sys.stdin):
part1 += run(blueprint, 24) * i
part2 *= run(blueprint, 32) if i<4 else 1
print(part1, part2)