adventofcode2022/day11/sol.py

72 lines
2.0 KiB
Python

from collections import Counter, namedtuple
import math
Monkey = namedtuple('Monkey', 'index, op, divisor, target')
def parse(input):
monkeys = []
items = {}
for i, chunk in enumerate(input.split("\n\n")):
lines = chunk.strip().split("\n")
assert 'Starting items:' in lines[1]
assert 'Operation: new =' in lines[2]
assert 'Test: divisible by' in lines[3]
assert 'If true: throw to' in lines[4]
assert 'If false: throw to' in lines[5]
items[i] = [int(x.strip()) for x in lines[1].split(":")[1].split(",")]
op = eval("lambda old: " + lines[2].split("=")[1])
divisor = int(lines[3].split()[-1])
target = [
int(lines[5].split()[-1]),
int(lines[4].split()[-1]),
]
monkeys.append(Monkey(items, op, divisor, target))
return monkeys, items
def play(monkeys, items, rounds=1, N=None):
throws = Counter()
for _ in range(rounds):
for i, m in enumerate(monkeys):
for j, x in enumerate(items[i]):
if N:
nx = m.op(x) % N
else:
nx = m.op(x) // 3
trg = m.target[nx % m.divisor == 0]
#print(f'{i}: {x} -> {nx}. throwing to {trg}')
items[trg].append(nx)
throws[i] += len(items[i])
items[i] = []
return throws
def show(items):
for k in items:
print(k, items[k])
def monkeybusiness(throws):
a, b = [x[1] for x in throws.most_common(2)]
return a*b
def lcm(ints):
n = 1
for x in ints:
n *= x // math.gcd(n, x)
return n
def main():
monkeys, items = parse(open("input").read())
#print(monkeys)
show(items)
tmp = {i: list(xs) for i, xs in items.items()}
throws = play(monkeys, tmp, rounds=20)
show(tmp)
print(monkeybusiness(throws))
N = lcm(m.divisor for m in monkeys)
throws = play(monkeys, items, rounds=10000, N=N)
show(items)
print(monkeybusiness(throws))
main()