aoc-2022/13.py

from itertools import zip_longest
from functools import cmp_to_key

def grouper(n, iterable, padvalue=None):
  return zip_longest(*[iter(iterable)]*n, fillvalue=padvalue)

lines = [eval(line.strip()) for line in open("13.input") if line.strip() != ""]

pairs = list(grouper(2, lines))

def comp(pair):
    left, right = pair

    if isinstance(left, int) and isinstance(right, int):
        if left == right:
            return None
        else:
            return left < right
    elif isinstance(left, list) and isinstance(right, list):
        for p in zip_longest(left, right, fillvalue=None):
            l, r = p
            if l is None:
                return True
            elif r is None:
                return False
            elif comp(p) is not None:
                return comp(p)
            else:
                continue
    else:
        if isinstance(left, int):
            return comp(([left], right))
        else:
            return comp((left, [right]))

s = 0

for i, pair in enumerate(pairs):
    if comp(pair):
        s += 1 + i

print("Part 1:", s)

packets = lines
packets.append([[2]])
packets.append([[6]])

packets.sort(key=cmp_to_key(lambda a, b: 1 if comp((a, b)) else -1), reverse=True)

prod = 1

for i, packet in enumerate(packets):
    if packet == [[2]] or packet == [[6]]:
        prod *= (1 + i)

print("Part 2:", prod)