import std/algorithm, std/os, std/sugar, std/strutils, std/math

proc read_crabs(filename: string): seq[int] =
   let f = open(filename)
   defer: f.close()
   result = collect:
      for crab in f.read_all().split(','):
         parse_int(crab.strip())

func part1(crabs: seq[int]): int =
   let median = sorted(crabs)[len(crabs) div 2]
   for crab in crabs:
      result += abs(crab - median)

func part2_fuel(crabs: seq[int], pos: int): int =
   for crab in crabs:
      let dist = abs(crab - pos)
      result += dist * (dist + 1) div 2

func part2(crabs: seq[int]): int =
   let
      l = len(crabs)
      s = sum(crabs)
   var c = 0
   for crab in crabs:
      if crab * l > s: inc c
   return part2_fuel(crabs, (s + c) div l)

when is_main_module:
   var filename: string = "7.in"
   if param_count() == 1:
      filename = param_str(1)
   let crabs = read_crabs(filename)
   echo part1(crabs)
   echo part2(crabs)