advent-of-code/2021/python/day07.py

38 lines
816 B
Python

import shared
from pprint import pp
from scipy import stats
import numpy as np
class Whale:
def __init__(self, name):
self.load(name)
self.set_target()
def load(self, name):
with open(name, "r") as f:
self.positions = list(map(int, f.read().split(",")))
def set_target(self):
data = np.array(self.positions)
self.target = round(np.mean(data))
def calculate(self, target=None):
return int(sum([self.sigma(abs(target - pos)) for pos in self.positions]))
def sigma(self, value):
return value * (value + 1) / 2
def main():
w = Whale(shared.get_fname(7))
if w.target in w.positions:
target = w.target
else:
target = w.target - 1
pp(w.calculate(target))
if __name__ == "__main__":
main()