aoc/2021/07/07.py

with open("input.txt") as f:
	crabs = [int(x) for x in f.read().split(',')]

from statistics import median, mean
from math import floor, ceil

def part1(crabs):
	m = int(median(crabs)) # the median is the midpoint
	offsets = [abs(x-m) for x in crabs] # that will be used to calculate simple cost.
	return sum(offsets)

print(part1(crabs))

def part2(crabs):
	def cost(d):
		return d*(d+1)//2 # sum of successive integers = triangle numbers.
	# because the mean is not an integer,
	m1 = floor(mean(crabs)) # we might need the floor,
	m2 = ceil(mean(crabs)) # or we might need the ceiling.
	fuel1 = [cost(abs(x-m1)) for x in crabs] # we calculate both costs
	fuel2 = [cost(abs(x-m2)) for x in crabs] # because it could be either one.
	return min([sum(fuel1), sum(fuel2)]) # just take the lesser of the two.

print(part2(crabs))