advent-of-code/2021/day7/day7_fast.py

30 lines
813 B
Python
Executable file

#!/usr/bin/env python
import sys
if __name__ == '__main__':
positions = [int(i) for i in sys.stdin.readline().split(',')]
def calc(starting_guess, distance_cost):
def target_cost(target):
return sum(distance_cost(abs(position-target)) for position in positions)
target = starting_guess(positions)
current_cost = target_cost(target)
for i in (-1, 1):
while (next_cost := target_cost(target+i)) < current_cost:
current_cost = next_cost
target += i
return current_cost
positions.sort()
print(calc(
lambda positions: positions[len(positions)//2],
lambda n: n
))
print(calc(
lambda positions: int(sum(positions)/len(positions)),
lambda n: n * (n+1)//2
))