Llm's idiomatic impl with groupby

day6/part2.py

@@ -1,31 +1,32 @@
+from itertools import groupby
 from math import prod
 from pathlib import Path
 
 FILE = "input.txt"
 
-def main():
-    raw_lines = Path(FILE).read_text().splitlines()
-    operators = raw_lines[-1].split()
-    operands: zip[tuple[str]] = zip(*raw_lines[:-1])
 
-    vals: list[int] = []
+def main(file: str = FILE) -> int:
-    sum = 0
+    *value_rows, operator_row = Path(file).read_text().splitlines()
-    for val in operands:
+    operators = operator_row.split()
-        if all(v == ' ' or v == '' for v in val):
+
-            sum += compute(vals, operators.pop(0))
+    # Columns represent vertical numbers; blank columns separate problems.
-            vals = []
+    columns = ("".join(col).strip() for col in zip(*value_rows))
-            continue
+    operands = (
-        vals.append(int("".join(val).strip()))
+        [int(n) for n in group]
-    return sum + compute(vals, operators.pop(0))
+        for is_value, group in groupby(columns, key=bool)
+        if is_value
+    )
+
+    return sum(compute(vals, op) for vals, op in zip(operands, operators))
+
 
 def compute(vals: list[int], op: str) -> int:
-    if op == '+':
+    ops = {"+": sum, "*": prod}
-        return sum(vals)
+    try:
-    elif op == '*':
+        return ops[op](vals)
-        return prod(vals)
+    except KeyError:
-    else:
+        raise ValueError(f"unknown operator: {op}")
-        print(f"unknown operator: {op}")
+
-        return 0
 
 if __name__ == "__main__":
     print(main())