Solve day 1 part 2

year_2023/day01/test_input_2.txt

@@ -0,0 +1,7 @@
+two1nine
+eightwothree
+abcone2threexyz
+xtwone3four
+4nineeightseven2
+zoneight234
+7pqrstsixteen

year_2023/day01/test_trebuchet.py

@@ -6,3 +6,10 @@
 
 def test_sum_calibration_values():
     assert sum_calibration_values(read_puzzle_input("test_input.txt")) == 142
+
+
+def test_sum_calibration_values_part_2():
+    assert (
+        sum_calibration_values(read_puzzle_input("test_input_2.txt"), parse_words=True)
+        == 281
+    )

year_2023/day01/trebuchet.py

@@ -1,10 +1,36 @@
 from util import read_puzzle_input
 
 
-def sum_calibration_values(puzzle_input):
+DIGITS = {
+    "one": "1",
+    "two": "2",
+    "three": "3",
+    "four": "4",
+    "five": "5",
+    "six": "6",
+    "seven": "7",
+    "eight": "8",
+    "nine": "9",
+}
+
+
+def _get_digit(value, idx):
+    if value[idx].isdigit():
+        return value[idx]
+    for length in [3, 4, 5]:
+        if value[idx : idx + length] in DIGITS:
+            return DIGITS[value[idx : idx + length]]
+
+
+def sum_calibration_values(puzzle_input, parse_words=False):
     answer = 0
     for value in puzzle_input.split("\n"):
-        digits = [c for c in value if c.isdigit()]
+        if parse_words:
+            digits = [
+                _get_digit(value, i) for i in range(len(value)) if _get_digit(value, i)
+            ]
+        else:
+            digits = [c for c in value if c.isdigit()]
         answer += int(digits[0] + digits[-1])
     return answer
 
@@ -13,4 +39,4 @@ def sum_calibration_values(puzzle_input):
     puzzle_input = read_puzzle_input()
 
     print(f"Part 1: {sum_calibration_values(puzzle_input)}")
-    print(f"Part 2: {sum_calibration_values(puzzle_input)}")
+    print(f"Part 2: {sum_calibration_values(puzzle_input, parse_words=True)}")