Day 18 Try 2

day_18.py

import copy
import functools
import typing as tp


Value = tp.Union[int, tp.List]
Pair = tp.List[Value]


def loader(lines: tp.Iterable[str]):
    for line in lines:
        if not line:
            break
        yield line


def parse(line: str) -> Pair:
    val, _pos = parse_rec(line, 0)
    return val


def parse_rec(line: str, pos: int) -> tp.Tuple[Pair, int]:
    assert line[pos] == "["
    pos += 1

    left_val, pos = parse_val(line, pos)

    assert line[pos] == ","
    pos += 1

    right_val, pos = parse_val(line, pos)

    assert line[pos] == "]"
    pos += 1

    return [left_val, right_val], pos


def parse_val(line, pos) -> tp.Tuple[Value, int]:
    if line[pos] == "[":
        val, pos = parse_rec(line, pos)
    else:
        val = str()
        while line[pos].isdigit():
            val += line[pos]
            pos += 1
        val = [int(val)]

    return val, pos


def explode(p: Pair) -> bool:
    return explode_rec(p, 0, [], [])


def descend(val: Value, index: int):
    while len(val) == 2:
        val = val[index]
    return val


def explode_rec(
    val: Value, level: int, left: tp.Optional[Value], right: tp.Optional[Value]
):
    if len(val) == 1:
        return False

    if level >= 4:
        if left:
            left[0] += val[0][0]
        if right:
            right[0] += val[1][0]
        val[:] = [0]
        return True

    return (
        explode_rec(val[0], level + 1, left, descend(val[1], 0))
        or
        explode_rec(val[1], level + 1, descend(val[0], 1), right)
    )


def split(val: Value):
    if len(val) == 1:
        return False
    return split_val(val, 0) or split_val(val, 1)


def split_val(val: Value, index: int):
    saved = val[index][0]
    if len(val[index]) == 1:
        if saved > 9:
            half = saved // 2
            val[index] = [[half], [saved - half]]
            return True
        return False
    return split(val[index])


def simplify(val):
    cond = True
    while cond:
        cond = False
        if explode(val):
            cond = True
            continue
        if split(val):
            cond = True
            continue
    return val


def add(a: Pair, b: Pair) -> Pair:
    a, b = copy.deepcopy(a), copy.deepcopy(b)
    return simplify([a, b])


def magnitude(val: Value):
    if len(val) == 1:
        return val[0]
    return 3 * magnitude(val[0]) + 2 * magnitude(val[1])


def task_1(lines: tp.Iterable[str]) -> int:
    snailfish_numbers = [parse(line) for line in loader(lines)]
    result = functools.reduce(add, snailfish_numbers)
    m = magnitude(result)
    return m


def task_2(lines: tp.Iterable[str]) -> int:
    snailfish_numbers = [parse(line) for line in loader(lines)]
    max_magnitude = 0
    for x in range(len(snailfish_numbers)):
        for y in range(len(snailfish_numbers)):
            if x == y:
                continue
            result = add(snailfish_numbers[x], snailfish_numbers[y])
            m = magnitude(result)
            max_magnitude = max(max_magnitude, m)

    return max_magnitude