aoc

Advent of Code
git clone https://git.pngu.org/aoc
Log | Files | Refs | README

19_2023.py (2935B)

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
import sys
import re
from copy import deepcopy
from math import gcd
from collections import defaultdict, Counter, deque
import heapq

D = open(sys.argv[1]).read().strip()
L = D.split("\n")
G = [[c for c in row] for row in L]
R = len(G)
C = len(G[0])

rules, parts = D.split("\n\n")
R = {}
for rule in rules.split("\n"):
    name, rest = rule.split("{")
    R[name] = rest[:-1]


def accepted(part):
    state = "in"
    while True:
        rule = R[state]
        for cmd in rule.split(","):
            applies = True
            res = cmd
            if ":" in cmd:
                cond, res = cmd.split(":")
                var = cond[0]
                op = cond[1]
                n = int(cond[2:])
                if op == ">":
                    applies = part[var] > n
                else:
                    applies = part[var] < n
            if applies:
                if res == "A":
                    return True
                if res == "R":
                    return False
                state = res
                break


ans = 0
for part in parts.split("\n"):
    part = part[1:-1]
    part = {x.split("=")[0]: int(x.split("=")[1]) for x in part.split(",")}
    if accepted(part):
        ans += part["x"] + part["m"] + part["a"] + part["s"]
print(ans)


def new_range(op, n, lo, hi):
    if op == ">":
        lo = max(lo, n + 1)
    elif op == "<":
        hi = min(hi, n - 1)
    elif op == ">=":
        lo = max(lo, n)
    elif op == "<=":
        hi = min(hi, n)
    else:
        assert False
    return (lo, hi)


def new_ranges(var, op, n, xl, xh, ml, mh, al, ah, sl, sh):
    if var == "x":
        xl, xh = new_range(op, n, xl, xh)
    elif var == "m":
        ml, mh = new_range(op, n, ml, mh)
    elif var == "a":
        al, ah = new_range(op, n, al, ah)
    elif var == "s":
        sl, sh = new_range(op, n, sl, sh)
    return (xl, xh, ml, mh, al, ah, sl, sh)


# x m a s
ans = 0
Q = deque([("in", 1, 4000, 1, 4000, 1, 4000, 1, 4000)])
while Q:
    state, xl, xh, ml, mh, al, ah, sl, sh = Q.pop()
    if xl > xh or ml > mh or al > ah or sl > sh:
        continue
    if state == "A":
        score = (xh - xl + 1) * (mh - ml + 1) * (ah - al + 1) * (sh - sl + 1)
        ans += score
        continue
    elif state == "R":
        continue
    else:
        rule = R[state]
        for cmd in rule.split(","):
            applies = True
            res = cmd
            if ":" in cmd:
                cond, res = cmd.split(":")
                var = cond[0]
                op = cond[1]
                n = int(cond[2:])
                Q.append((res, *new_ranges(var, op, n, xl, xh, ml, mh, al, ah, sl, sh)))
                xl, xh, ml, mh, al, ah, sl, sh = new_ranges(
                    var, "<=" if op == ">" else ">=", n, xl, xh, ml, mh, al, ah, sl, sh
                )
            else:
                Q.append((res, xl, xh, ml, mh, al, ah, sl, sh))
                break
print(ans)