path: root/day7.py

#!/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)