Today’s problem was a kind-of-a-tree-but-not-really adventure, and was the first of this year’s problems where a naïve, brute-force implemention didn’t work for part 1.
I didn’t have much time today, and I’m not really happy with my solution. Especially two things stand out:
- The
stepfunctions for the folds are really similar between part 1 and part 2, but I couldn’t find a quick way to abstract over the differences. I feel like there is some common structure here that I should be able to use. - I’m using unsafe indexing into arrays. I should be able to tell Lean about the
dimensions, possibly by using a
Nat-indexed vector type.
Things I (re-)learned today
- Nothing, really. Did a quick’n’dirty Haskell-in-Lean solution today. :|
Solution
import Batteries.Data.List
import Aoc
open Aoc
abbrev Set := Std.HashSet
def Part1.step (beams : Set Nat) (line : List Char) : Nat × (Set Nat) :=
line.zipIdx.foldl (λ acc@⟨hits, beams⟩ (c, xPos) =>
match c with
| '^' => if xPos ∈ beams then
let beams' := beams.erase xPos
|>.insert (xPos - 1)
|>.insert (xPos + 1)
(hits + 1, beams')
else
acc
| _ => acc
) (0, beams)
def Part2.step (acc : Array Nat × Set Nat) (line : List Char) : Array Nat × Set Nat :=
let (timelines, beams) := acc
line.zipIdx.foldl (λ acc@⟨timelines, beams⟩ (c, xPos) =>
match c with
| '^' =>
if xPos ∈ beams then
let beams' := beams.erase xPos
|>.insert (xPos - 1)
|>.insert (xPos + 1)
let timelinesAtXPos := timelines[xPos]!
let timelines' := timelines
-- the timeline at the current x-position ended …
|>.set! xPos 0
-- … and was split into two new timelines
|>.modify (xPos-1) (timelinesAtXPos + ·)
|>.modify (xPos+1) (timelinesAtXPos + ·)
(timelines', beams')
else
acc
| _ => acc
) (timelines, beams)
def main : IO Unit := do
let lines ← readLines "Day07/input.txt"
let first := lines[0]!
let width := first.length
let startPos := first.toList.idxOf 'S' |> (Position.mk · 0)
-- part 1
let (numSplitters, _) := lines.map (·.toList) |>.foldl (λ acc line =>
let (newHits, beams') := Part1.step acc.snd line
(acc.fst + newHits, beams')
) (0, Std.HashSet.emptyWithCapacity.insert startPos.x.toNat)
IO.println s!"Part 1: {numSplitters}"
-- part 2
-- initially, there is a single timeline at the start position
let initialTimelines : Array Nat :=
Array.replicate width 0 |>.set! startPos.x.toNat 1
let initialBeams : Set Nat := Std.HashSet.ofList [startPos.x.toNat]
let (numTimelines, _) := lines.map (·.toList)
|>.foldl Part2.step (initialTimelines, initialBeams)
IO.println s!"Part 2: {numTimelines.sum}"