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))