139 lines
3.9 KiB
Python
Executable file
139 lines
3.9 KiB
Python
Executable file
#!/usr/bin/env python
|
|
|
|
import enum
|
|
import functools
|
|
import itertools
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, Iterator, List, Tuple, TypeVar
|
|
|
|
RawPacket = List[bool]
|
|
|
|
|
|
class PacketType(enum.IntEnum):
|
|
SUM = 0
|
|
PRODUCT = 1
|
|
MINIMUM = 2
|
|
MAXIMUM = 3
|
|
LITTERAL = 4
|
|
GREATER = 5
|
|
LESS = 6
|
|
EQUAL = 7
|
|
|
|
|
|
class PacketLengthType(enum.IntEnum):
|
|
TOTAL_BITS = 0 # Next 15 bits are the number of bits in the sub-packets
|
|
TOTAL_PACKETS = 1 # Next 11 bits are the number of sub-packets
|
|
|
|
|
|
@dataclass
|
|
class PacketLength:
|
|
type: PacketLengthType
|
|
length: int
|
|
|
|
|
|
@dataclass
|
|
class Packet:
|
|
version: int
|
|
type: PacketType
|
|
|
|
|
|
@dataclass
|
|
class OperatorPacket(Packet):
|
|
packets: List[Packet]
|
|
|
|
|
|
@dataclass
|
|
class LitteralPacket(Packet):
|
|
number: int
|
|
|
|
|
|
def solve(input: List[str]) -> int:
|
|
def to_raw(packet: str) -> RawPacket:
|
|
def bits(n: int) -> Iterator[bool]:
|
|
for i in range(3, -1, -1):
|
|
yield bool(n & (1 << i))
|
|
|
|
nums = [int(c, 16) for c in input[0]]
|
|
return list(itertools.chain.from_iterable(bits(n) for n in nums))
|
|
|
|
def bits_to_int(bits: Iterable[bool]) -> int:
|
|
return functools.reduce(lambda a, b: (a << 1) + b, bits, 0)
|
|
|
|
def parse_packet(p: RawPacket) -> Tuple[int, Packet]:
|
|
def packet_version(p: RawPacket) -> int:
|
|
return bits_to_int(p[:3])
|
|
|
|
def packet_type(p: RawPacket) -> PacketType:
|
|
return PacketType(bits_to_int(p[3:6]))
|
|
|
|
def parse_length(p: RawPacket) -> Tuple[int, PacketLength]:
|
|
assert packet_type(p) != PacketType.LITTERAL # Sanity check
|
|
type = PacketLengthType(bits_to_int(p[6:7]))
|
|
if type == PacketLengthType.TOTAL_BITS:
|
|
index = 7 + 15
|
|
length = bits_to_int(p[7:index])
|
|
else:
|
|
index = 7 + 11
|
|
length = bits_to_int(p[7:index])
|
|
return index, PacketLength(type, length)
|
|
|
|
def parse_litteral(p: RawPacket) -> Tuple[int, LitteralPacket]:
|
|
version, type = packet_version(p), packet_type(p)
|
|
assert type == PacketType.LITTERAL # Sanity check
|
|
index = 6
|
|
bits: List[bool] = []
|
|
while True:
|
|
bits += p[index + 1 : index + 5]
|
|
index += 5
|
|
# Check if we were at the last one
|
|
if p[index - 5] == 0:
|
|
break
|
|
return index, LitteralPacket(version, type, bits_to_int(bits))
|
|
|
|
def parse_operator(p: RawPacket) -> Tuple[int, OperatorPacket]:
|
|
version, type = packet_version(p), packet_type(p)
|
|
assert type != PacketType.LITTERAL # Sanity check
|
|
|
|
index, length = parse_length(p)
|
|
packets: List[Packet] = []
|
|
|
|
if length.type == PacketLengthType.TOTAL_BITS:
|
|
sub_index = 0
|
|
while sub_index < length.length:
|
|
parsed, packet = parse_packet(p[index:])
|
|
sub_index += parsed
|
|
index += parsed
|
|
packets.append(packet)
|
|
else:
|
|
while len(packets) < length.length:
|
|
parsed, packet = parse_packet(p[index:])
|
|
index += parsed
|
|
packets.append(packet)
|
|
|
|
return index, OperatorPacket(version, type, packets)
|
|
|
|
if packet_type(p) == PacketType.LITTERAL:
|
|
return parse_litteral(p)
|
|
return parse_operator(p)
|
|
|
|
def score(p: Packet) -> int:
|
|
if p.type == PacketType.LITTERAL:
|
|
return p.version
|
|
assert isinstance(p, OperatorPacket) # Sanity check
|
|
return p.version + sum(score(c) for c in p.packets)
|
|
|
|
raw = to_raw(input[0])
|
|
__, packet = parse_packet(raw)
|
|
|
|
return score(packet)
|
|
|
|
|
|
def main() -> None:
|
|
input = [line.strip() for line in sys.stdin.readlines()]
|
|
print(solve(input))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|