AOC2021 Day 18 (Python)

day18_tests.py

import functools
import math


def value_at(smallfish_number, path):
    cursor = smallfish_number
    for e in path:
        if e == 'L':
            cursor = cursor[0]
        else:
            cursor = cursor[1]
    return cursor


def get_indexed_values(smallfish_number, path_so_far):
    l, r = smallfish_number
    if type(l) is list:
        yield from get_indexed_values(l, path_so_far + 'L')
    else:
        yield path_so_far + 'L', l

    if type(r) is list:
        yield from get_indexed_values(r, path_so_far + 'R')
    else:
        yield path_so_far + 'R', r


def update(smallfish_number, path, new_value):
    if len(path) == 0:
        return new_value
    l, r = smallfish_number
    if path[0] == 'L':
        return [update(l, path[1:], new_value), r]
    else:
        return [l, update(r, path[1:], new_value)]


def explode(smallfish_number, before_indexed_value, indexed_value, after_indexed_value):
    path, _ = indexed_value
    containing_path = path[:-1]
    l, r = value_at(smallfish_number, containing_path)

    if before_indexed_value:
        before_path, before_value = before_indexed_value
        smallfish_number = update(smallfish_number, before_path, before_value + l)

    if after_indexed_value:
        after_path, after_value = after_indexed_value
        smallfish_number = update(smallfish_number, after_path, after_value + r)

    smallfish_number = update(smallfish_number, containing_path, 0)
    return smallfish_number


def reduce(smallfish_number):

    while True:
        has_changed = False

        indexed_values = list(get_indexed_values(smallfish_number, ''))
        for index, (path, value) in enumerate(indexed_values):
            if len(path) == 5:
                smallfish_number = explode(
                    smallfish_number,
                    indexed_values[index - 1] if index > 0 else None,
                    (path, value),
                    indexed_values[index + 2] if index < len(indexed_values) - 2 else None)
                has_changed = True
                break

        if has_changed:
            continue

        for path, value in indexed_values:
            if value >= 10:
                smallfish_number = update(smallfish_number,
                                          path,
                                          [math.floor(value / 2.0), math.ceil(value / 2.0)])
                has_changed = True
                break

        if not has_changed:
            break

    return smallfish_number


def test_reduce():
    print()
    assert reduce([[[[[4, 3], 4], 4], [7, [[8, 4], 9]]], [1, 1]]) == \
           [[[[0, 7], 4], [[7, 8], [6, 0]]], [8, 1]]


def read_smallfish_number(line):
    cursor = []
    stack = []
    for c in line:
        if c == '[':
            stack.append(cursor)
            cursor = []
        elif c == ']':
            parent = stack.pop()
            parent.append(cursor)
            cursor = parent
        elif str.isdigit(c):
            cursor.append(int(c))
    return cursor[0]


def test_read_smallfish_number():
    assert read_smallfish_number('[[[[[4, 3], 4], 4], [7, [[8, 4], 9]]], [1, 1]]') == \
           [[[[[4, 3], 4], 4], [7, [[8, 4], 9]]], [1, 1]]


def add_smallfish_numbers(l, r):
    result = reduce([l, r])
    return result


def magnitude(smallfish_number):
    l, r = smallfish_number
    l_mag = magnitude(l) if type(l) is list else l
    r_mag = magnitude(r) if type(r) is list else r
    return (3 * l_mag) + (2 * r_mag)


def test_add_smallfish_numbers():
    example1 = """
[[[0,[4,5]],[0,0]],[[[4,5],[2,6]],[9,5]]]
[7,[[[3,7],[4,3]],[[6,3],[8,8]]]]
[[2,[[0,8],[3,4]]],[[[6,7],1],[7,[1,6]]]]
[[[[2,4],7],[6,[0,5]]],[[[6,8],[2,8]],[[2,1],[4,5]]]]
[7,[5,[[3,8],[1,4]]]]
[[2,[2,2]],[8,[8,1]]]
[2,9]
[1,[[[9,3],9],[[9,0],[0,7]]]]
[[[5,[7,4]],7],1]
[[[[4,2],2],6],[8,7]]
    """.strip().split('\n')

    print()
    added = functools.reduce(add_smallfish_numbers, (read_smallfish_number(line) for line in example1))
    assert added == [[[[8, 7], [7, 7]], [[8, 6], [7, 7]]], [[[0, 7], [6, 6]], [8, 7]]]

    example2 = """
[[[0,[5,8]],[[1,7],[9,6]]],[[4,[1,2]],[[1,4],2]]]
[[[5,[2,8]],4],[5,[[9,9],0]]]
[6,[[[6,2],[5,6]],[[7,6],[4,7]]]]
[[[6,[0,7]],[0,9]],[4,[9,[9,0]]]]
[[[7,[6,4]],[3,[1,3]]],[[[5,5],1],9]]
[[6,[[7,3],[3,2]]],[[[3,8],[5,7]],4]]
[[[[5,4],[7,7]],8],[[8,3],8]]
[[9,3],[[9,9],[6,[4,9]]]]
[[2,[[7,7],7]],[[5,8],[[9,3],[0,2]]]]
[[[[5,2],5],[8,[3,7]]],[[5,[7,5]],[4,4]]]
""".strip().split('\n')

    added = functools.reduce(add_smallfish_numbers, (read_smallfish_number(line) for line in example2))
    assert added == [[[[6, 6], [7, 6]], [[7, 7], [7, 0]]], [[[7, 7], [7, 7]], [[7, 8], [9, 9]]]]

    with open('puzzle_inputs/day18.txt') as file:
        numbers = [read_smallfish_number(line) for line in file]

    added = functools.reduce(add_smallfish_numbers, numbers)
    print()
    print(magnitude(added))

    print(max(magnitude(add_smallfish_numbers(a, b)) for a in numbers for b in numbers if a != b))