aoc/2023/19/solution.py

#!/usr/bin/python3
import sys
import re

def less(a, b):
    return a < b

def greater(a, b):
    return a > b

def part1():
    print('part 1')
    flow = {}
    pos = {'x': 0, 'm': 1, 'a': 2, 's': 3}
    for line in map(str.strip, sys.stdin):
        if len(line) == 0:
            break
        workflow, rulestring = re.findall(r'([a-z]+)\{(.*)\}', line)[0]
        rules = rulestring.split(',')
        flow[workflow] = []
        for rule in rules:
            extracted = re.findall(r'([ARa-z]+)?([><])?([0-9]+)?:?([ARa-z]+)?', rule)[0]
            if extracted[1] == '':
                flow[workflow].append((extracted[0],))
            else:
                flow[workflow].append((pos[extracted[0]], less if extracted[1] == '<' else greater, int(extracted[2]), extracted[3]))
    ratings = 0
    for line in map(str.strip, sys.stdin):
        stats = [int(n) for n in re.findall(r'\{x=(\d+),m=(\d+),a=(\d+),s=(\d+)\}', line)[0]]
        workflow = 'in'
        while workflow not in 'RA':
            rules = flow[workflow]
            for rule in rules:
                if len(rule) == 1:
                    workflow = rule[0]
                    break
                else:
                    i, fun, value, target = rule
                    if fun(stats[i], value):
                        workflow = target
                        break
            if workflow == 'A':
                ratings += sum(stats)
    print(ratings)

def part2():
    print('part 2')
    flow = {}
    pos = {'x': 0, 'm': 1, 'a': 2, 's': 3}
    for line in map(str.strip, sys.stdin):
        if len(line) == 0:
            break
        workflow, rulestring = re.findall(r'([a-z]+)\{(.*)\}', line)[0]
        rules = rulestring.split(',')
        flow[workflow] = []
        for rule in rules:
            extracted = re.findall(r'([ARa-z]+)?([><])?([0-9]+)?:?([ARa-z]+)?', rule)[0]
            if extracted[1] == '':
                flow[workflow].append((extracted[0],))
            else:
                flow[workflow].append((pos[extracted[0]], less if extracted[1] == '<' else greater, int(extracted[2]), extracted[3]))
    ratings = 0
    workflow = 'in'
    ranges = [(1, 4000), (1, 4000), (1, 4000), (1, 4000)]
    while workflow not in 'RA':
        ranges, workflow = nodes.pop()
        rules = flow[workflow]
        for rule in rules:
            if len(rule) == 1:
                workflow = rule[0]
                break
            else:
                i, fun, value, target = rule
                if fun(stats[i], value):
                    workflow = target
                    break
        if workflow == 'A':
            ratings += sum(stats)
    print(ratings)

if sys.argv[1] in '1':
    part1()
else:
    part2()