Advent of Code day 16 part 2

AoC_2021_16_2.py

"""
My solution for AoC 2021 day 16 part 2.
Most of it was taken from @benediktwerner's solution:
https://github.com/benediktwerner/AdventOfCode/blob/master/2021/day16/sol.py
"""
from math import prod

OP = {
    0: sum,
    1: prod,
    2: min,
    3: max,
    5: lambda lst: int(lst[0] > lst[1]),
    6: lambda lst: int(lst[0] < lst[1]),
    7: lambda lst: int(lst[0] == lst[1]),
}


def calculate(data):

    ptype = int(data[3:6], 2)

    if ptype == 4:  # data is literal value packet
        bits, data = "", data[6:]
        while data[0] == "1":
            bits += data[1:5]
            data = data[5:]
        return int(bits + data[1:5], 2), data[5:]

    else:  # data is packet of operator type
        args = []
        if data[6] == "0":
            length = int(data[7:22], 2)
            end = 22 + length
            rest = data[22:end]
            while rest:
                arg, rest = calculate(rest)
                args.append(arg)
            data = data[end:]
        else:
            count = int(data[7:18], 2)
            data = data[18:]
            for _ in range(count):
                arg, data = calculate(data)
                args.append(arg)
        return OP[ptype](args), data


hexstr = open("input_16.txt").read().strip()
full_packet = "".join(f"{int(ch, 16):04b}" for ch in hexstr)
result, _ = calculate(full_packet)
print(result)