01.lean

def Input := List Int × List Int

def parse (input : String) : Input :=
  input
    |>.trim
    |>.splitOn "\n"
    |>.map (
      fun line => line
        |>.splitOn " "
        |>.filter (not ∘ String.isEmpty)
        |>.map String.toInt!
        |> fun nums => (nums.head!, nums.getLast!)
      )
    |>.unzip

def part1 (input : Input) : Nat :=
  let fst := input.fst.mergeSort (· ≤ ·)
  let snd := input.snd.mergeSort (· ≤ ·)
  fst
    |>.zip snd
    |>.map (fun (a, b) => a - b |>.natAbs)
    |>.foldl Nat.add 0

def part2 (input : Input) : Int :=
  input.fst
    |>.map (fun a => a * (input.snd.count a))
    |>.foldl Int.add 0

def main (args : List String) : IO Unit := do
  let input ← parse <$> (System.FilePath.mk args[0]! |> IO.FS.readFile)
  IO.println $ part1 input
  IO.println $ part2 input