19_parts.py 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. #!/usr/bin/env python3
  2. import sys
  3. import re
  4. def parse_rule(rule):
  5. if ':' not in rule:
  6. return 0, (1, 4000), rule
  7. left, op, right, _, destination = re.split('([<>:])', rule)
  8. rating = 'xmas'.index(left)
  9. rating_range = (1, int(right) - 1) if op == '<' else (int(right) + 1, 4000)
  10. return rating, rating_range, destination
  11. def parse(inp):
  12. flows = {}
  13. for line in inp:
  14. if line == '\n':
  15. break
  16. flow, rule_strings = line[:-2].split('{')
  17. flows[flow] = list(map(parse_rule, rule_strings.split(',')))
  18. parts = [tuple(map(int, re.findall(r'\d+', line))) for line in inp]
  19. return flows, parts
  20. def at(ranges, i, lower, upper):
  21. return *ranges[:i], (lower, upper), *ranges[i + 1:]
  22. def naccept(ranges, flows, flow):
  23. if flow == 'R':
  24. return 0
  25. if flow == 'A':
  26. x, m, a, s = (upper - lower + 1 for lower, upper in ranges)
  27. return x * m * a * s
  28. acc = 0
  29. for i, (imin, imax), dst in flows[flow]:
  30. rmin, rmax = ranges[i]
  31. lower = max(rmin, imin)
  32. upper = min(rmax, imax)
  33. if upper >= lower:
  34. acc += naccept(at(ranges, i, lower, upper), flows, dst)
  35. if rmin < imin:
  36. acc += naccept(at(ranges, i, rmin, imin - 1), flows, flow)
  37. if rmax > imax:
  38. acc += naccept(at(ranges, i, imax + 1, rmax), flows, flow)
  39. break
  40. return acc
  41. flows, parts = parse(sys.stdin)
  42. print(sum(sum(p) for p in parts if naccept(tuple(zip(p, p)), flows, 'in')))
  43. print(naccept(((1, 4000),) * 4, flows, 'in'))