Today’s problem was quite fun, but I got myself really confused in part 2. Nothing Lean-specific today, unfortunately, as I spent way too long getting it to work at all.

My first solution for part 1, was quite simple: Just keep track of the largest predecessor digit:

def Part1.findMaxJoltage (batteryBank: List Nat) : Nat :=
  batteryBank.foldl (fun ⟨largestPred, largest⟩ jolt =>
      let joltVal := largestPred * 10 + jolt
      (Nat.max largestPred jolt, Nat.max largest joltVal)
    ) ((0, 0) : Nat × Nat)
    |> (·.snd)

Then, for part 2, where I had to generalize it to n batteries. I had an idea that I could work from right to left and consider each digit and do something like the following:

• if larger than current first digit: include it and turn off the lowest digit in the current number string
• if equal to current first digit: include, and if it was larger than any existing digit, then turn off that digit
• if less than current first digit: skip

Spoiler alert: This will never work. However, it did work for the examples in the problem description, so it took a while before I realized I had to abandon this method and start over again and instead work from left to right and “push” each digit past any smaller digits, as far left as I could, while ensuring that there were enough digits left to complete the n-digit number.

Since I was working with numbers represented as a list of digits today, I added a small helper to my utility library:

def List.digitsToNat (ds : List Nat) : Nat :=
  ds.foldl (· * 10 + ·) 0

Things I (re-)learned today

• While working on a misguided solution and looking for a tails function, remembered that I can use the Batteries package, so I was able to remove the scanl function I added to my utility library on day 1 and use the one in Batteries.Data.List instead.
• I can use assert! to check invariants during development.
• I have to use let rec to define local, recursive functions, like in SML/OCaml (if I remember correctly).

Solution

import Aoc
open Aoc

def findMaxJoltage (numBatteries : Nat) (batteryBank: List Nat) :=
  let rec popWhileSmaller stack currentDigit remaining :=
    match stack with
    | [] => []
    | d :: ds =>
      if d < currentDigit && (ds.length + remaining >= numBatteries)
        then popWhileSmaller ds currentDigit remaining
        else stack

  batteryBank.foldl (fun ⟨stack, numRemaining⟩ currentDigit =>
      let stack := popWhileSmaller stack currentDigit numRemaining
      let stack := if stack.length < numBatteries
                     then currentDigit :: stack
                     else stack
      (stack, numRemaining - 1)
    ) ([], batteryBank.length) |> (·.fst.reverse.digitsToNat)

def main : IO Unit := do
  let input ← Aoc.readNonEmptyLines "Day03/input.txt"
  let batteryBanks := Array.map (·.toList.map (·.toNat - '0'.toNat)) input

  -- Part 1
  let maxJoltages := batteryBanks.map (findMaxJoltage 2)
  IO.println s!"Part 1: {maxJoltages.sum}"

  -- Part 2
  let maxJoltages := batteryBanks.map (findMaxJoltage 12)
  IO.println s!"Part 2: {maxJoltages.sum}"