aocpy/2021/d16.py
2024-11-23 14:54:02 -05:00

86 lines
1.6 KiB
Python

from lib import get_data
data = get_data(__file__)
bs = ""
for c in data.strip():
b = bin(int(c, 16))[2:]
b = "0" * (4 - len(b)) + b
bs += b
total_v = 0
def parse(bs: str, i: int):
v = int(bs[i : i + 3], 2)
global total_v
total_v += v
i += 3
t = int(bs[i : i + 3], 2)
i += 3
# print(v, t)
if t == 4:
lv = ""
while True:
seg = bs[i : i + 5]
lv += seg[1:]
i += 5
if seg[0] == "0":
break
lv = int(lv, 2)
return i, lv
vs = []
if bs[i] == "0":
# total length of bits
i += 1
bits = int(bs[i : i + 15], 2)
i += 15
i_max = i + bits
while True:
i, v = parse(bs, i)
vs.append(v)
if i >= i_max:
break
elif bs[i] == "1":
# number of sub-packets
i += 1
num_sub = int(bs[i : i + 11], 2)
i += 11
vs = []
for _ in range(num_sub):
i, v = parse(bs, i)
vs.append(v)
else:
assert False
if t == 0:
return i, sum(vs)
elif t == 1:
x = 1
for v in vs:
x *= v
return i, x
elif t == 2:
return i, min(vs)
elif t == 3:
return i, max(vs)
elif t == 5:
assert len(vs) == 2
return i, 1 if vs[0] > vs[1] else 0
elif t == 6:
assert len(vs) == 2
return i, 1 if vs[0] < vs[1] else 0
elif t == 7:
assert len(vs) == 2
return i, 1 if vs[0] == vs[1] else 0
else:
assert False
p2 = parse(bs, 0)[1]
print(total_v)
print(p2)