aoc-2020/18.py

lines = []

with open("18.input") as f:
    for line in f.readlines():
        lines.append(line.strip())

def parse(line):
    res = []

    i = 0
    while i < len(line):
        c = line[i]
        i += 1
        if c in "0123456789+*":
            res.append(c)
        elif c == '(':
            buf = []
            depth = 1
            while True:
                c = line[i]
                i += 1
                if c == '(':
                    depth += 1
                elif c == ')':
                    depth -= 1
                if depth == 0:
                    break
                buf.append(c)
            res.append(parse(buf))
    return res


def task_1(expression):
    res = int(expression[0])
    i = 1
    while i < len(expression):
        op = expression[i]
        i += 1
        val = int(expression[i])
        i += 1
        if op == '*':
            res *= val
        elif op == '+':
            res += val
    return res


def task_2(expression):
    # evaluate additions
    i = 0
    while i < len(expression):
        if expression[i] == '+':
            expression[i - 1] = int(expression[i - 1]) + int(expression[i + 1])
            del expression[i:i+2]
            i = i - 1
        i += 1
    # evaluate multiplications
    i = 0
    while i < len(expression):
        if expression[i] == '*':
            expression[i - 1] = int(expression[i - 1]) * int(expression[i + 1])
            del expression[i:i+2]
            i = i - 1
        i += 1
    assert len(expression) == 1
    return expression[0]


def evaluate(expression, presedence_function):
    if len(expression) == 1:
        try:
            return int(expression[0])
        except ValueError:
            return evaluate(expression[0], presedence_function)
    else:
        expression = [e if type(e) is not list else evaluate(e, presedence_function) for e in expression]
        # now we no parentheses left :-)
        return presedence_function(expression)


print("Answer 1:", sum([evaluate(parse(line), task_1) for line in lines]))
print("Answer 2:", sum([evaluate(parse(line), task_2) for line in lines]))