aoc2023/d19.py

from lib import *

EXAMPLE = """px{a<2006:qkq,m>2090:A,rfg}
pv{a>1716:R,A}
lnx{m>1548:A,A}
rfg{s<537:gd,x>2440:R,A}
qs{s>3448:A,lnx}
qkq{x<1416:A,crn}
crn{x>2662:A,R}
in{s<1351:px,qqz}
qqz{s>2770:qs,m<1801:hdj,R}
gd{a>3333:R,R}
hdj{m>838:A,pv}

{x=787,m=2655,a=1222,s=2876}
{x=1679,m=44,a=2067,s=496}
{x=2036,m=264,a=79,s=2244}
{x=2461,m=1339,a=466,s=291}
{x=2127,m=1623,a=2188,s=1013}
"""

def solve(i: Input, second=False):
    res = 0
    ps = i.paras()
    wfs, parts = ps
    wfs = wfs.splitlines()
    parts = parts.splitlines()

    workflows = {}
    for w in wfs:
        name, ins = w.split("{")
        ins = ins.replace("}", "")
        conds = ins.split(",")
        wf = []
        for cond in conds[:-1]:
            cmp, n = cond.split(":")
            if '<' in cmp:
                a, b = cmp.split('<')
                wf.append((a, 'smaller', int(b), n))
            elif '>' in cmp:
                a, b = cmp.split('>')
                wf.append((a, 'greater', int(b), n))
            else:
                raise Exception()
        wf.append(conds[-1:])
        workflows[name] = wf

    def parse_part(part):
        d = {}
        p = part.replace("{", "").replace("}", "")
        for pair in p.split(","):
            a, b = pair.split("=")
            d[a] = int(b)
        return d

    parts = list(map(parse_part, parts))
    if not second:
        for p in parts:
            current = 'in'
            while current not in ['A', 'R']:
                for inst in workflows[current]:
                    if len(inst) == 4:
                        letter = inst[0]
                        value = inst[2]
                        next = inst[3]
                        if inst[1] == 'smaller':
                            if p[letter] < value:
                                current = next
                                break
                        elif inst[1] == 'greater':
                            if p[letter] > value:
                                current = next
                                break
                        else:
                            raise Exception()
                    elif len(inst) == 1:
                        current = inst[0]
                    else:
                        raise Exception()
            if current == 'A':
                r = sum(p.values())
                res += r
        return res

    ranges = [{c: (1, 4000) for c in 'xmas'}]
    ranges[0]['cur'] = 'in'
    ranges[0]['idx'] = 0
    accepted = []
    while ranges:
        r = ranges.pop()
        cur, idx = r['cur'], r['idx']

        if cur == 'A':
            accepted.append(r)
            continue
        elif cur == 'R':
            continue

        inst = workflows[cur][idx]
        if len(inst) == 4:
            letter = inst[0]
            value = inst[2]
            nxt = inst[3]
            ro = r[letter]
            r1, r2 = dict(r), dict(r)
            if inst[1] == 'smaller':
                r1[letter] = (ro[0], value - 1)
                r1['idx'] = 0
                r1['cur'] = nxt
                r2[letter] = (value, ro[1])
                r2['idx'] += 1
            elif inst[1] == 'greater':
                r1[letter] = (ro[0], value)
                r1['idx'] += 1
                r2[letter] = (value + 1, ro[1])
                r2['idx'] = 0
                r2['cur'] = nxt
            if r1[letter][1] >= r1[letter][0]:
                ranges.append(r1)
            if r2[letter][1] >= r2[letter][0]:
                ranges.append(r2)
        elif len(inst) == 1:
            r['cur'] = inst[0]
            r['idx'] = 0
            ranges.append(r)
    res = 0
    for a in accepted:
        r = 1
        for c in 'xmas':
            l, u = a[c]
            r *= (u + 1 - l)
        res += r
    return res

def main():
    DAY_INPUT = "i19.txt"

    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    # 25:00

    print("Example  2:", solve(Input(EXAMPLE), True))
    assert solve(Input(EXAMPLE), True) == 167409079868000
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    # 120:00

if __name__ == "__main__":
    main()