import re from functools import reduce test_input_p1 = ( "xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))" ) test_input_p2 = ( "xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))" ) test_solution_p1 = 161 test_solution_p2 = 48 instruction_pattern = re.compile(r"mul\((\d+),(\d+)\)|(do(?:n't)?)\(\)") def solve_p1(puzzle_input: str) -> int: instructions = ( tuple(map(int, match[:2])) for match in instruction_pattern.findall(puzzle_input) if match[0] ) return sum(a * b for a, b in instructions) def solve_p2(puzzle_input: str) -> int: instructions = ( match[2] or tuple(map(int, match[:2])) for match in instruction_pattern.findall(puzzle_input) ) return reduce( lambda a, i: ( (1 if i == "do" else 0, a[1]) if isinstance(i, str) else (a[0], a[1] + a[0] * i[0] * i[1]) ), instructions, (1, 0), )[1]