Day 07, Part 2

scompt · scompt · commit a13d77048ec6 · 2021-12-07T08:06:25.000+01:00
diff --git a/day_07/__main__.py b/day_07/__main__.py
@@ -5,16 +5,20 @@ def read_input(textio: TextIO) -> List[int]:
     positions = [int(num) for num in textio.readlines()[0].split(',')]
     return positions
 
-def cost(positions:List[int], target_position:int) -> int:
-    return sum(abs(position-target_position) for position in positions)
+def cost(position: int, target_position: int) -> int:
+    diff = abs(position-target_position)
+    return int((diff+1)*diff/2)
+
+def total_cost(positions:List[int], target_position:int) -> int:
+    return sum(cost(position, target_position) for position in positions)
 
 def find_cheapest(positions:List[int]):
     min_pos = min(positions)
     max_pos = max(positions)
 
     costs = dict()
     for pos in range(min_pos, max_pos+1):
-        costs[pos] = cost(positions, pos)
+        costs[pos] = total_cost(positions, pos)
 
     min_cost_pos = min(costs, key=costs.get)
     return min_cost_pos, costs[min_cost_pos]
