33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
|
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)
|