refactor for geneeral n_batteries case

Author: jameslawlor

src/advent_of_code/year_2025/day_03.py

@@ -8,7 +8,7 @@ def combine_joltages(list_of_joltages):
     return int("".join([str(x) for x in list_of_joltages]))
 
 
-def calculate_largest_joltage(battery_bank_as_ints):
+def calculate_largest_joltage(battery_bank_as_ints, n_batteries):
     largest_joltage_1 = 0
 
     for battery_index, battery_value in enumerate(battery_bank_as_ints):
@@ -22,31 +22,19 @@ def calculate_largest_joltage(battery_bank_as_ints):
     return combine_joltages([largest_joltage_1, largest_joltage_2])
 
 
-def calculate_largest_joltage_part_2(battery_bank_as_ints, n_batteries=12):
-    sorted_batteries = sorted(battery_bank_as_ints, reverse=True)
-    return combine_joltages(sorted_batteries[:n_batteries])
-
-
-def solve_part(input, part=1):
+def solve_part(input, n_batteries):
     joltage_list = []
 
-    if part == 1:
-        calculate_largest_joltage_function = calculate_largest_joltage
-    elif part == 2:
-        calculate_largest_joltage_function = calculate_largest_joltage_part_2
-    else:
-        raise ValueError("Part must be 1 or 2")
-    
     for battery_bank_str in input:
         battery_bank_as_ints = parse_battery_bank_to_ints_list(battery_bank_str)
-        largest_joltage = calculate_largest_joltage_function(battery_bank_as_ints)
+        largest_joltage = calculate_largest_joltage(battery_bank_as_ints, n_batteries)
         joltage_list.append(largest_joltage)
 
     return sum(joltage_list)
 
 def solve(input):
-    part_1_solution = solve_part(input, part=1)
-    part_2_solution = solve_part(input, part=2)
+    part_1_solution = solve_part(input, n_batteries=2)
+    part_2_solution = solve_part(input, n_batteries=12)
     return (part_1_solution, part_2_solution)
 
 

tests/year_2025/test_day_03.py

@@ -5,7 +5,6 @@
     calculate_largest_joltage,
     parse_battery_bank_to_ints_list,
     combine_joltages,
-    calculate_largest_joltage_part_2,
 )
 
 
@@ -29,7 +28,6 @@ def test_solver(day_03_test_input, day_03_expected_output):
     assert result == day_03_expected_output
 
 
-
 @pytest.fixture
 def test_calculate_largest_joltage_params():
     return [
@@ -39,17 +37,11 @@ def test_calculate_largest_joltage_params():
         ("818181911112111", 92, 888911112111),
     ]
 
-def test_calculate_largest_joltage_part_1(test_calculate_largest_joltage_params):
-    for input_str, expected_output, _ in test_calculate_largest_joltage_params:
+def test_calculate_largest_joltage(test_calculate_largest_joltage_params):
+    for input_str, expected_output_part_1, expected_output_part_2 in test_calculate_largest_joltage_params:
         parsed_inputs = parse_battery_bank_to_ints_list(input_str)
-        assert expected_output == calculate_largest_joltage(parsed_inputs)
-
-
-def test_calculate_largest_joltage_part_2(test_calculate_largest_joltage_params):
-    for input_str, _, expected_output in test_calculate_largest_joltage_params:
-        parsed_inputs = parse_battery_bank_to_ints_list(input_str)
-        assert expected_output == calculate_largest_joltage_part_2(parsed_inputs)
-
+        assert expected_output_part_1 == calculate_largest_joltage(parsed_inputs, n_batteries=2)
+        assert expected_output_part_2 == calculate_largest_joltage(parsed_inputs, n_batteries=12)
 
 def test_parse_battery_bank_to_ints_list():
     assert parse_battery_bank_to_ints_list("98765") == [9,8,7,6,5]