Solve 2021 day 16

2021/d16.py

@@ -0,0 +1,85 @@
+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)