#!/usr/bin/env python import numpy as np def triangle(num): return (num * (num + 1)) / 2 with open("day7.txt") as data: distances = np.array([num for num in map(int, next(data).strip().split(','))]).astype(int) # part 1 median = np.median(distances) print(np.sum(np.abs(distances - median).astype(int))) # part 2 minim, maxim = np.amin(distances), np.amax(distances) least_dist = None for i in range(minim, maxim + 1): dist = np.sum(triangle(np.abs(distances - i)).astype(int)) if least_dist is None or least_dist > dist: least_dist = dist print(least_dist)