51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
import re
|
|
from collections import namedtuple
|
|
from typing import Generator, List, TextIO
|
|
|
|
|
|
Mul = namedtuple("Mul", ["x", "y"])
|
|
MulPattern = re.compile(r"mul\((\d{1,3}),(\d{1,3})\)")
|
|
ExtendedPattern = re.compile(r"(?P<op_mul>mul)\((\d{1,3}),(\d{1,3})\)|(?P<op_do>do)\(\)|(?P<op_dont>don't)\(\)")
|
|
|
|
|
|
def parse_mul_ops(f: TextIO) -> List[Mul]:
|
|
input_string = f.read()
|
|
return [Mul(int(x), int(y)) for x, y in MulPattern.findall(input_string)]
|
|
|
|
def parse_mul_ops_extended(f: TextIO) -> Generator[Mul, None, None]:
|
|
input_string = f.read()
|
|
|
|
mul_enabled = True
|
|
for match in ExtendedPattern.finditer(input_string):
|
|
|
|
if match.group("op_dont"):
|
|
mul_enabled = False
|
|
elif match.group("op_do"):
|
|
mul_enabled = True
|
|
|
|
if match.group("op_mul") and mul_enabled:
|
|
yield Mul(int(match.group(2)), int(match.group(3)))
|
|
|
|
|
|
def main():
|
|
input_file = "input.txt"
|
|
|
|
total = 0
|
|
total_extended = 0
|
|
with open(input_file, "r") as f:
|
|
for x, y in parse_mul_ops(f):
|
|
print(f"mul({x}, {y}) = {x * y}")
|
|
total += x * y
|
|
|
|
with open(input_file, "r") as f:
|
|
for x, y in parse_mul_ops_extended(f):
|
|
print(f"mul({x}, {y}) = {x * y}")
|
|
total_extended += x * y
|
|
|
|
print(f"total = {total}")
|
|
print(f"total_extended = {total_extended}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|