aoc-2021/18.py

from util import get_input
from math import floor, ceil
from functools import reduce
from itertools import product

def parse(line):
    return eval(line)

input = get_input("18.input", parse)

def add(a, b):
    return red([a, b])

def red(a_in):
    a = list(a_in)
    while True:
        old_a = list(a)
        a = try_explode(a, 0)[0]
        if a != old_a:
            continue
        a = try_split(a)
        if old_a == a:
            break
    return a


def add_const(a, c, left):
    if type(a) == int:
        return a + c
    elif left:
        return [add_const(a[0], c, left), a[1]]
    else:
        return [a[0], add_const(a[1], c, left)]


def try_explode(a, depth):
    if type(a) == int:
        return a, [0, 0], False
    if depth == 4:
        return 0, a, True
    a0, [l, r], ex = try_explode(a[0], depth + 1)
    if ex:
        return [a0, add_const(a[1], r, True)], [l, 0], True
    a1, [l, r], ex = try_explode(a[1], depth + 1)
    if ex:
        return [add_const(a[0], l, False), a1], [0, r], True
    return [a0, a1], [0, 0], False

def test_explode():
    print(try_explode([[[[[9,8],1],2],3],4], 0)[0])
    print(try_explode([7,[6,[5,[4,[3,2]]]]], 0)[0])
    print(try_explode([[6,[5,[4,[3,2]]]],1], 0)[0])
    print(try_explode([[3,[2,[1,[7,3]]]],[6,[5,[4,[3,2]]]]], 0)[0])
    print(try_explode([[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]], 0)[0])

def try_split(a):
    if type(a) == int:
        if a >= 10:
            return [floor(a / 2), ceil(a / 2)]
        else:
            return a
    else:
        l = try_split(a[0])
        if l != a[0]:
            return [l, a[1]]
        else:
            return [a[0], try_split(a[1])]

def test_add():
    print(add([[[[4,3],4],4],[7,[[8,4],9]]], [1,1]))
    t2 = [[1,1],
            [2,2],
            [3,3],
            [4,4]]
    print(reduce(add, t2))
    t3 = [[1,1],
            [2,2],
            [3,3],
            [4,4],
            [5,5],
            [6,6]]
    print(reduce(add, t3))

def magnitude(a):
    if type(a) == int:
        return a
    else:
        return 3 * magnitude(a[0]) + 2 * magnitude(a[1])

print("Part 1:", magnitude(reduce(add, input)))

print("Part 2:", max(magnitude(add(a, b)) for [a, b] in product(input, repeat=2)))