Improve 2024 day 3 regex

2024/d3.py

@@ -5,22 +5,20 @@ from lib import ints
 
 data = get_data(__file__)
 
-t = 0
-r = re.compile(r'mul\(\d+,\d+\)')
-for match in r.findall(data):
-    a, b = ints(match)
-    t += a * b
-print(t)
-
-t = 0
+t1 = 0
+t2 = 0
 enabled = True
-for i in range(len(data.strip())):
-    m = r.match(data[i:])
-    if m and enabled:
-        a, b = ints(m.group())
-        t += a * b
-    elif data[i:].startswith("do()"):
+r = re.compile(r"mul\(\d+,\d+\)|do\(\)|don't\(\)")
+for m in r.findall(data):
+    if m == "do()":
         enabled = True
-    elif data[i:].startswith("don't()"):
+    elif m == "don't()":
         enabled = False
-print(t)
+    else:
+        a, b = ints(m)
+        t1 += a * b
+        if enabled:
+            t2 += a * b
+
+print(t1)
+print(t2)