68 lines
1.6 KiB
Python
68 lines
1.6 KiB
Python
from typing import Iterator
|
|
import sys
|
|
import re
|
|
|
|
mul_re = re.compile(
|
|
r"mul\((?P<n1>\d+),(?P<n2>\d+)\)"
|
|
)
|
|
|
|
mul_re_cond = re.compile(
|
|
r"mul\((?P<n1>\d+),(?P<n2>\d+)\)|(?P<do>do\(\))|(?P<dont>don't\(\))"
|
|
)
|
|
|
|
def process_basic(data) -> int:
|
|
total = 0
|
|
for match in re.finditer(mul_re, data):
|
|
n1 = int(match.group('n1'))
|
|
n2 = int(match.group('n2'))
|
|
total += n1 * n2
|
|
return total
|
|
|
|
def process_cond(data):
|
|
# Could also split into segments between do and the next don't
|
|
# And process each block between them
|
|
# But this is simpler
|
|
total = 0
|
|
enabled = True
|
|
for match in re.finditer(mul_re_cond, data):
|
|
if match.group('do'):
|
|
enabled = True
|
|
elif match.group('dont'):
|
|
enabled = False
|
|
elif enabled:
|
|
n1 = int(match.group('n1'))
|
|
n2 = int(match.group('n2'))
|
|
total += n1 * n2
|
|
return total
|
|
|
|
|
|
def load_data(filename) -> str:
|
|
with open(filename) as f:
|
|
return f.read()
|
|
|
|
def test():
|
|
print("Beginning Test 1")
|
|
data = r"xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))"
|
|
assert process_basic(data) == 161
|
|
print("Test 1 passed")
|
|
|
|
print("Beginning Test 2")
|
|
data = r"xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))"
|
|
assert process_cond(data) == 48
|
|
print("Test 2 passed")
|
|
|
|
def main():
|
|
test()
|
|
data = load_data(sys.argv[1])
|
|
total = process_basic(data)
|
|
print(f"Total from data without conditionals: {total}")
|
|
total = process_cond(data)
|
|
print(f"Total from data with conditionals: {total}")
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|