Advent of code 2024 day 5

def day5(s, part2=False):
  p1, p2 = s.split('\n\n')
  rules = {tuple(line.split('|')) for line in p1.splitlines()}
  updates = (line.split(',') for line in p2.splitlines())

  def compare(a, b):
    return -1 if (a, b) in rules else 1 if (b, a) in rules else 0

  total = 0
  for update in updates:
    new = sorted(update, key=functools.cmp_to_key(compare))
    if (new == update) ^ part2:
      total += int(new[len(new) // 2])

  return total