advent-of-code/2022/day13.py

42 lines
1.3 KiB
Python
Raw Permalink Normal View History

2022-12-13 07:54:15 +00:00
import json
import sys
from functools import cmp_to_key
2023-12-09 08:54:53 +00:00
from typing import TypeAlias, cast
2022-12-13 07:54:15 +00:00
blocks = sys.stdin.read().strip().split("\n\n")
pairs = [tuple(json.loads(p) for p in block.split("\n")) for block in blocks]
2023-12-09 08:54:53 +00:00
Packet: TypeAlias = list[int | list["Packet"]]
2022-12-13 07:54:15 +00:00
2023-12-09 08:54:53 +00:00
def compare(lhs: Packet, rhs: Packet) -> int:
2022-12-13 07:59:53 +00:00
for lhs_a, rhs_a in zip(lhs, rhs):
if isinstance(lhs_a, int) and isinstance(rhs_a, int):
if lhs_a != rhs_a:
return rhs_a - lhs_a
2022-12-13 07:54:15 +00:00
else:
2022-12-13 07:59:53 +00:00
if not isinstance(lhs_a, list):
2023-12-09 08:54:53 +00:00
lhs_a = [lhs_a] # type: ignore
2022-12-13 07:59:53 +00:00
elif not isinstance(rhs_a, list):
2023-12-09 08:54:53 +00:00
rhs_a = [rhs_a] # type: ignore
2022-12-13 07:59:53 +00:00
assert isinstance(rhs_a, list) and isinstance(lhs_a, list)
2023-12-09 08:54:53 +00:00
r = compare(cast(Packet, lhs_a), cast(Packet, rhs_a))
2022-12-13 07:59:53 +00:00
if r != 0:
return r
return len(rhs) - len(lhs)
answer_1 = sum(i + 1 for i, (lhs, rhs) in enumerate(pairs) if compare(lhs, rhs) > 0)
2022-12-13 07:54:15 +00:00
print(f"answer_1 is {answer_1}")
dividers = [[[2]], [[6]]]
packets = [packet for packets in pairs for packet in packets]
packets.extend(dividers)
packets = list(reversed(sorted(packets, key=cmp_to_key(compare))))
d_index = [packets.index(d) + 1 for d in dividers]
print(f"answer 2 is {d_index[0] * d_index[1]}")