adventofcode2022/day19/sol.py

import math
import re

sample = {
    1:{
        'ore': [4, 0, 0, 0],
        'clay': [2, 0, 0, 0],
        'obsidian': [3, 14, 0, 0],
        'geode': [2, 0, 7, 0],
    },
    2: {
        'ore': [2, 0, 0, 0],
        'clay': [3, 0, 0, 0],
        'obsidian': [3, 8, 0, 0],
        'geode': [3, 0, 12, 0],
    }
}

robot_number = {'ore': 0, 'clay': 1, 'obsidian': 2, 'geode': 3}

def parse(line):
    line = re.sub(r'[^\d]+', ' ', line)
    idx, ore1, ore2, ore3, clay3, ore4, obs4 = map(int, line.split())
    return idx, {
        'ore': [ore1, 0,0,0],
        'clay': [ore2, 0,0,0],
        'obsidian': [ore3,clay3,0,0],
        'geode': [ore4,0,obs4,0],
    }

def simulate(blueprint, minutes=24):
    rmax = [max(x[i] for x in blueprint.values()) for i in range(3)] + [9999]

    limit = None
    if minutes >= 32:
        limit = 100000

    items = list(blueprint.items())
    items.reverse()

    print(items)

    buckets = [[] for _ in range(minutes+1)]
    def enqueue(t, robots, resources):
        assert t > 0
        if t <= 0 or t >= len(buckets):
            return
        x = resources[3]
        buckets[t].append((x,robots, resources))

    enqueue(minutes, robots=[1,0,0,0], resources=[0,0,0,0])

    max_geodes = 0

    for _ in range(minutes):
        buckets[minutes].sort(reverse=True)
        #print(minutes, buckets[minutes][:1])
        for _, robots, resources in buckets[minutes][:limit]:
            if 1:
                geodes = robots[3]*minutes + resources[3]
                if geodes > max_geodes:
                    max_geodes = geodes
                elif geodes < max_geodes:
                    # if the number of geodes we could get if we stopped building robots
                    # is less than the max such number we've seen so far,
                    # then prune this node.
                    #
                    # this seems like it would be too greedy a rule, but
                    # somehow it actually works. (i think this is equivalant to
                    # the rule "always build a geode robot as soon as you can"
                    # and its validity depends on the input data. (in fact it
                    # *doesn't* work for the first sample blueprint for part 2.
                    # but it works on my input so w/e)
                    #
                    # decreases the search space enormously
                    continue
            for robot, cost in items:
                i = robot_number[robot]
                if robots[i] >= rmax[i]:
                    # don't build more of 1 robot than we can spend in 1 minute
                    continue
                if robot != 'geode' and robots[i]*minutes + resources[i] >= minutes*rmax[i]:
                    # insight from reddit: if we have enough resources on hand and enough
                    # mining capacity to build any robot that needs it every turn until
                    # time is up, then we don't need to build any more of that kind of robot
                    #
                    # this provides a minor speed-up over just doing the simpler check
                    continue

                # figure out how soon we can afford it
                wait = 0
                for x,y,r in zip(resources, cost, robots):
                    if y == 0:
                        continue
                    if r <= 0:
                        wait = 9999
                        break
                    if y > x:
                        wait = max(wait, int(math.ceil((y - x)/r)))
                if wait+1 >= minutes:
                    continue

                new_resources = [x+(wait+1)*r-y for x,y,r in zip(resources, cost, robots)]
                new_robots = list(robots)
                new_robots[i] += 1
                enqueue(minutes-wait-1, new_robots, new_resources)

        print(minutes, max_geodes, [len(x) for x in buckets[:minutes+1]])
        buckets[minutes] = [] # clear old buckets to save memory

        minutes -= 1

    return max_geodes

input = {}
with open('input') as f:
    for line in f:
        idx, bp = parse(line)
        input[idx] = bp

def solve(input):
    t = 0
    for idx, B in input.items():
        g = simulate(B)
        t += idx*g
        print("blueprint", idx, "max geodes is", g)
        print("---")
    print(t)
    return t

def solve2(input):
    t = 1
    for idx, B in input.items():
        if idx <= 3:
            g = simulate(B, minutes=32)
            t *= g
    print(t)
    return t

assert solve(sample) == 33
solve(input)
solve2(input)