Day 18 Try 1

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:
    _a, _b, exp = explode_rec(p, 0)
    return exp


def explode_rec(val: Value, level: int):
    if level == 4:
        return val[0], val[1], True

    if isinstance(val[0], list):
        a, b, exp = explode_rec(val[0], level + 1)
        if exp:
            if b is not None:
                if isinstance(val[1], int):
                    val[1] += b
                else:
                    add_val(val[1], b, 0)
                b = None
                if level == 3:
                    val[0] = 0
            return a, b, exp

    if isinstance(val[1], list):
        a, b, exp = explode_rec(val[1], level + 1)
        if exp:
            if a is not None:
                if isinstance(val[0], int):
                    val[0] += a
                else:
                    add_val(val[0], a, 1)
                a = None
                if level == 3:
                    val[1] = 0
            return a, b, exp

    return None, None, False


def explode_rec2(val: Value, level: int):
    if isinstance(val[0], list):
        if level == 3:
            a, b = val[0]
            val[0] = 0
            if isinstance(val[1], int):
                val[1] += b
            else:
                add_val(val[1], b, 0)
            return a, None, True
        else:
            a, b, exp = explode_rec2(val[0], level + 1)
            if exp:
                return a, b, exp

    if isinstance(val[1], list):
        if level == 3:
            a, b = val[1]
            val[1] = 0
            if isinstance(val[0], int):
                val[0] += a
            else:
                add_val(val[0], a, 1)
            return None, b, True
        else:
            a, b, exp = explode_rec2(val[1], level + 1)
            if exp:
                return a, b, exp

    return None, None, False


def add_val(val: Value, number: int, index: int):
    if isinstance(val[index], int):
        val[index] += number
    else:
        add_val(val[index], number, index)


def split(val: Value):
    if isinstance(val, int):
        return False
    return split_val(val, 0) or split_val(val, 1)


def split_val(val: Value, index: int):
    saved = val[index]
    if isinstance(saved, int) and saved > 9:
        val[index] = [saved // 2, saved - saved // 2]
        return True
    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 isinstance(val, int):
        return val
    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