11_monkeys.py 1.2 KB

1234567891011121314151617181920212223242526272829303132
  1. #!/usr/bin/env python3
  2. import sys
  3. from operator import add, mul, pow
  4. from functools import reduce
  5. from heapq import nlargest
  6. def parse(f):
  7. for line in f:
  8. if line.startswith(' Starting items:'):
  9. items = tuple(map(int, line.split(':')[1].split(', ')))
  10. opstr, amount = next(f).replace('* old', '** 2').split()[-2:]
  11. op = {'+': add, '*': mul, '**': pow}[opstr], int(amount)
  12. throws = tuple(int(next(f).rsplit(' ', 1)[1]) for _ in range(3))
  13. yield items, op, throws
  14. def business(monkeys, rounds, divisor):
  15. queues = [list(items) for items, _, _ in monkeys]
  16. inspects = [0] * len(monkeys)
  17. modulus = reduce(mul, (div for _, _, (div, _, _) in monkeys))
  18. for _ in range(rounds):
  19. for i, (_, (op, amount), (div, true, false)) in enumerate(monkeys):
  20. items = queues[i]
  21. for item in items:
  22. worry = op(item, amount) // divisor % modulus
  23. queues[false if worry % div else true].append(worry)
  24. inspects[i] += len(items)
  25. items.clear()
  26. return mul(*nlargest(2, inspects))
  27. monkeys = list(parse(sys.stdin))
  28. print(business(monkeys, 20, 3))
  29. print(business(monkeys, 10000, 1))