To get the optimal position for the crabs, we just need to move them all to the median of their starting positions. Python makes this super-simple, just use staticstics.median().
from statistics import median
def optimal_fuel_cost(positions: list[int]) -> int:
target = int(median(positions))
return sum(abs(pos - target) for pos in positions)
test_positions = [16, 1, 2, 0, 4, 2, 7, 1, 2, 14]
assert optimal_fuel_cost(test_positions) == 37
import aocd
positions = [int(pos) for pos in aocd.get_data(day=7, year=2021).split(",")]
print("Part 1:", optimal_fuel_cost(positions))
Part 1: 339321
The new best position is the (rounded) mean of their starting positions, and the fuel cost for a given crab to move to the new position can be calculated using the triangle number formula.
That's because we are trying to minimize the sum over the triangle distances from crab positions $c$ to target position $t$ (you can ignore the division by two, it doesn't matter to the outcome):
$$min_t \sum_{i=1}^n \frac {(|c_i - t|) (|c_i - t| + 1)} {2}$$and $(|c_i - t|) (|c_i - t| + 1)$ is very close to $(c_i - t) ^ 2$. When solve for $t$ we get:
$$t = \frac {\sum_{i=1}^n c_i} {n} + \frac {\sum_{i=1}^n sgn(c_i - t)} {2n}$$Here, the $sgn()$ function gives -1, 0 or 1 for negative, zero or positive outcomes of $c_i - t$. The first term in the sum is just the mean of the positions (the sum of all distances divided by $n$, the number of values). The other part is always going to be somewhere between -0.5 and +0.5 (the value will always lie between $\frac {-n} {2n}$ and $\frac {n} {2n}$)!
So the trick lies in where to round off to; I simply round up and down (add -0.5, 0 or 0.5 and use int() to floor the value), and see which one is cheaper.
from statistics import mean
def optimal_fuel_cost_triangle(positions: list[int]) -> int:
target = mean(positions)
def triangle(n: int) -> int:
return n * (n + 1) // 2
def summed_cost(t: int) -> int:
return sum(triangle(abs(pos - t)) for pos in positions)
return min(summed_cost(int(target + bias)) for bias in (-0.5, 0, 0.5))
assert optimal_fuel_cost_triangle(test_positions) == 168
print("Part 2:", optimal_fuel_cost_triangle(positions))
Part 2: 95476244