adventofcode-2021/day07/day07.py

#!/usr/bin/env python3
import sys
from typing import Sequence


def rangesum(rn: range):
    assert rn.step == 1, "this function doesn't work when range step != 1"
    return (((rn.stop + 1) * rn.stop) // 2) - (((rn.start + 1) * rn.start) // 2)


def part1(positions: Sequence[int]):
    lo = min(positions)
    hi = max(positions)
    diffs = []
    for i in range(lo, hi + 1):
        diffs += [sum(abs(pos - i) for pos in positions)]
    print(min(diffs))


def part2(positions: Sequence[int]):
    lo = min(positions)
    hi = max(positions)
    ranges = []
    for i in range(lo, hi + 1):
        ranges += [[range(abs(pos - i)) for pos in positions]]
    diffs = [sum(rangesum(r) for r in rs) for rs in ranges]
    print(min(diffs))


positions = list(map(int, sys.stdin.readline().strip().split(",")))

print("Part 1")
part1(positions)
print("Part 2")
part2(positions)