advent-of-code-2024/day7/code.py

import sys
sys.path.append('../aoclib')

from aoclib import Input

def parse_input(lines: list):
    equations = list()
    for l in lines:
        s = l.split(':')
        eq = (int(s[0]), [int(x.strip()) for x in s[1].split()])
        equations.append(eq)    
    return equations

def _sum2(l: list[int]) -> list[int]:
    return [l[0] + l[1], *l[2:]]

def _mul2(l: list[int]) -> list[int]:
    return [l[0] * l[1], *l[2:]]

def _concat2(l: list[int]) -> list[int]:
    return [int(str(l[0]) + str(l[1])), *l[2:]]

def may_be_true(val: int, lst: list[int], funcs) -> bool:
    if len(lst) == 1:
        return lst[0] == val
    return any([may_be_true(val, f(lst), funcs) for f in funcs])

def solve1(eqs):
    funcs = [_sum2, _mul2]
    return sum([e[0] if may_be_true(e[0], e[1], funcs) else 0 for e in eqs]) 

def solve2(eqs):
    funcs = [_sum2, _mul2, _concat2]
    return sum([e[0] if may_be_true(e[0], e[1], funcs) else 0 for e in eqs]) 
    

if __name__ == '__main__':
    #lines = Input('input_test.txt').lines()
    lines = Input('input.txt').lines()

    eq = parse_input(lines)

    #part 1
    print('part 1:', solve1(eq))

    #part 2
    print('part 2:', solve2(eq))