import sys import numpy as np positions = np.asarray([int(c) for c in sys.stdin.read().strip().split(",")]) min_position, max_position = positions.min(), positions.max() # part 1 answer_1 = min( np.sum(np.abs(positions - position)) for position in range(min_position, max_position + 1) ) print(f"answer 1 is {answer_1}") # part 2 answer_2 = min( np.sum(abs(positions - position) * (abs(positions - position) + 1) // 2) for position in range(min_position, max_position + 1) ) print(f"answer 2 is {answer_2}")