Advent of Code 2024 - Day 4

Stream: general

[James Sully]

Here's my solution. Feedback appreciated, and please post your own solution!

import Utils
namespace Day04

abbrev CharGrid := Array (Array Char)

def parse (input : String) : CharGrid :=
  let lines := input.lines.toArray
  lines.map (·.toList.toArray)

structure V2 (α : Type) where
  x : α
  y : α

def V2.map (f : α → β) (v : V2 α) : V2 β := ⟨f v.x, f v.y⟩

def V2.mapM [Monad m] (f : α → m β) (u : V2 α) : m (V2 β) := do
  let x ← f u.x
  let y ← f u.y
  return {x, y}

instance [HAdd α β γ] : HAdd (V2 α) (V2 β) (V2 γ) where
  hAdd | ⟨u₁, u₂⟩, ⟨v₁, v₂⟩ => ⟨u₁+v₁, u₂+v₂⟩

instance [HMul α β γ] : HMul α (V2 β) (V2 γ) where
  hMul | a, ⟨x, y⟩ => ⟨a*x, a*y⟩

def gridLookup (idx : V2 Nat) (grid : Array (Array α)) : Option α := do
  let row ← grid[idx.y]?
  let x ← row[idx.x]?
  return x

def gridFindAll [Inhabited α] (p : α → Bool) (arr : Array (Array α))
  : Array (V2 Nat) := Id.run do
  let height := arr.size
  let width := arr[0]!.size
  let mut ret := #[]
  for row in [:height] do
    for col in [:width] do
      if p arr[row]![col]! then
        ret := ret.push {x := col, y := row}
  ret

def directions : List (V2 Int) := [
  ⟨0, 1⟩,   -- N
  ⟨1, 1⟩,   -- NE
  ⟨1, 0⟩,   -- E
  ⟨1, -1⟩,  -- SE
  ⟨0, -1⟩,  -- S
  ⟨-1, -1⟩, -- SW
  ⟨-1, 0⟩,  -- W
  ⟨-1, 1⟩   -- NW
]

def xmas := "XMAS".toList

def isXmasInDir (grid : CharGrid) (start : V2 Int) (direction : V2 Int) : Bool :=
  let coords : List (V2 Int) := List.range 4 |>.map (λ i ↦ start + Int.ofNat i * direction)
  match do
    let coords : List (V2 Nat) ← coords.mapM (V2.mapM Int.toNat')
    let chars : List Char ← coords.mapM (gridLookup · grid)
    return chars
    with
  | .none => false
  | .some chars => chars = xmas

def countXmassStartingAt (grid : CharGrid) (start : V2 Int) : Nat :=
  directions.countP (isXmasInDir grid start ·)

def part1 (input : CharGrid) : Nat := Id.run do
  let xPositions : Array (V2 Nat) := gridFindAll (· = 'X') input
  xPositions.map (·.map Int.ofNat)
  |>.map (countXmassStartingAt input ·)
  |>.toList
  |>.sum

def mas := "MAS".toList

-- this converting between Int and Nat coords sucks
-- we'll use 'A' as the "origin"
def isCrossMasAt (grid : CharGrid) (aPosition : V2 Nat) : Bool :=
  let aPosition' : V2 Int := aPosition.map Int.ofNat
  match do
    let diag1Deltas : List (V2 Int) := [⟨-1,-1⟩, ⟨0,0⟩, ⟨1, 1⟩]
    let diag1Coords ← diag1Deltas.map (aPosition' + ·) |>.mapM (V2.mapM Int.toNat')
    let diag1 ← diag1Coords.mapM (gridLookup · grid)

    let diag2Deltas : List (V2 Int) := [⟨-1,1⟩, ⟨0,0⟩, ⟨1, -1⟩]
    let diag2Coords ← diag2Deltas.map (aPosition' + ·) |>.mapM (V2.mapM Int.toNat')
    let diag2 ← diag2Coords.mapM (gridLookup · grid)

    return (diag1, diag2)
    with
    | .none => false
    | .some (diag1, diag2) =>
      (diag1 = mas || diag1.reverse = mas) && (diag2 = mas || diag2.reverse = mas)

def part2 (input : CharGrid) : Nat :=
  let aPositions := gridFindAll (· = 'A') input
  aPositions.toList.countP (isCrossMasAt input ·)

[Kim Morrison]

def gridFindAll [Inhabited α] (p : α → Bool) (arr : Array (Array α)) :
    Array (V2 Nat) := Id.run do
  let height := arr.size
  let mut ret := #[]
  for h₁ : row in [:height] do
    for h₂ : col in [:arr[row].size] do
      if p arr[row][col] then
        ret := ret.push {x := col, y := row}
  ret

[Kim Morrison]

You could also do that in a more functional style with List.findIdxs from Batteries (although there is no Array analogue at present).

[Kim Morrison]

def gridLookup (idx : V2 Nat) (grid : Array (Array α)) : Option α := do
  let row ← grid[idx.y]?
  let x ← row[idx.x]?
  return x

does seem to be calling out for syntactical sugar to allow grid[idx.x, idx.y]?

[Kim Morrison]

Could you quarantine the Int/Nat pain by changing gridLookup to accept Ints?

[James Sully]

that gridFindAll is great, thanks, also eliminates the Inhabited constraint

[James Sully]

I want to try and avoid batteries, I added it for the HashMap, then removed it after you pointed out the Std one hahaha. It's nice to avoid dependencies for AOC.

[James Sully]

Kim Morrison said:

Could you quarantine the Int/Nat pain by changing gridLookup to accept Ints?

I think this would probably be the practical choice yeah. Feels immoral since all negatives will be none

[James Sully]

Kim Morrison said:

grid[idx.x, idx.y]?

that would be rad

[James Sully]

I think the "correct" solution for the Int/Nat thing would be a fallible add V2 Nat -> V2 Int -> Option (V2 Nat)
But then you'd have to give up + or make your own notation.

[James Sully]

wait could I do a HAdd instance for that?

[Kim Morrison]

James Sully said:

Kim Morrison said:

grid[idx.x, idx.y]?

that would be rad

This is something someone should prototype in a downstream library. It shouldn't need any native language support to implement, I think.

[Matt Russell]

Day 4 solution from a Lean beginner, feedback / suggestions welcome!
https://github.com/mdr/aoc-2024/blob/main/Aoc2024/Day04/Solve.lean

[James Sully]

The Int/Nat situation improved quite a bit with a HAdd instance:

import Utils
namespace Day04

abbrev CharGrid := Array (Array Char)

def parse (input : String) : CharGrid :=
  let lines := input.lines.toArray
  lines.map (·.toList.toArray)

structure V2 (α : Type) where
  x : α
  y : α

def V2.map (f : α → β) (v : V2 α) : V2 β := ⟨f v.x, f v.y⟩

def V2.mapM [Monad m] (f : α → m β) (u : V2 α) : m (V2 β) := do
  let x ← f u.x
  let y ← f u.y
  return {x, y}

instance [HAdd α β γ] : HAdd (V2 α) (V2 β) (V2 γ) where
  hAdd | ⟨u₁, u₂⟩, ⟨v₁, v₂⟩ => ⟨u₁+v₁, u₂+v₂⟩

instance : HAdd (V2 Nat) (V2 Int) (Option (V2 Nat)) where
  hAdd | ⟨u₁, u₂⟩, ⟨v₁, v₂⟩ => do
    let x ← Int.ofNat u₁ + v₁ |>.toNat'
    let y ← Int.ofNat u₂ + v₂ |>.toNat'
    return ⟨x, y⟩

instance [HMul α β γ] : HMul α (V2 β) (V2 γ) where
  hMul | a, ⟨x, y⟩ => ⟨a*x, a*y⟩

def gridLookup (idx : V2 Nat) (grid : Array (Array α)) : Option α := do
  let row ← grid[idx.y]?
  let x ← row[idx.x]?
  return x

def gridFindAll (p : α → Bool) (arr : Array (Array α)) :
    Array (V2 Nat) := Id.run do
  let height := arr.size
  let mut ret := #[]
  for h₁ : row in [:height] do
    for h₂ : col in [:arr[row].size] do
      if p arr[row][col] then
        ret := ret.push {x := col, y := row}
  ret

def directions : List (V2 Int) := [
  ⟨0, 1⟩,   -- N
  ⟨1, 1⟩,   -- NE
  ⟨1, 0⟩,   -- E
  ⟨1, -1⟩,  -- SE
  ⟨0, -1⟩,  -- S
  ⟨-1, -1⟩, -- SW
  ⟨-1, 0⟩,  -- W
  ⟨-1, 1⟩   -- NW
]

def xmas := "XMAS".toList

def isXmasInDir (grid : CharGrid) (start : V2 Nat) (direction : V2 Int) : Bool :=
  match do
    let coords : List (V2 Nat) ← List.range 4 |>.mapM (λ i ↦ start + Int.ofNat i * direction)
    let chars : List Char ← coords.mapM (gridLookup · grid)
    return chars
    with
  | .none => false
  | .some chars => chars = xmas

def countXmassStartingAt (grid : CharGrid) (start : V2 Nat) : Nat :=
  directions.countP (isXmasInDir grid start ·)

def part1 (input : CharGrid) : Nat :=
  let xPositions : Array (V2 Nat) := gridFindAll (· = 'X') input
  xPositions
  |>.map (countXmassStartingAt input ·)
  |>.toList
  |>.sum

def mas := "MAS".toList

-- we'll use 'A' as the "origin"
def isCrossMasAt (grid : CharGrid) (aPosition : V2 Nat) : Bool :=
  match do
    let diag1Deltas : List (V2 Int) := [⟨-1,-1⟩, ⟨0,0⟩, ⟨1, 1⟩]
    let diag1Coords : List (V2 Nat) ← diag1Deltas.mapM (λ (d : V2 Int) ↦ aPosition + d)
    let diag1 ← diag1Coords.mapM (gridLookup · grid)

    let diag2Deltas : List (V2 Int) := [⟨-1,1⟩, ⟨0,0⟩, ⟨1, -1⟩]
    let diag2Coords : List (V2 Nat) ← diag2Deltas.mapM (aPosition + ·)
    let diag2 ← diag2Coords.mapM (gridLookup · grid)

    .some (diag1, diag2)
    with
    | .none => false
    | .some (diag1, diag2) =>
      (diag1 = mas || diag1.reverse = mas) && (diag2 = mas || diag2.reverse = mas)

def part2 (input : CharGrid) : Nat :=
  let aPositions := gridFindAll (· = 'A') input
  aPositions.toList.countP (isCrossMasAt input ·)

I'm not entirely sure why I need the explicit lambda with a type annotation on the argument on this line

    let diag1Coords : List (V2 Nat) ← diag1Deltas.mapM (λ (d : V2 Int) ↦ aPosition + d)

but the type checker gets upset without it.

I think this is probably a little too magical to really be nice code. Could be good with a semantic separation between point and vector types.

[James Sully]

Matt Russell said:

Day 4 solution from a Lean beginner, feedback / suggestions welcome!
https://github.com/mdr/aoc-2024/blob/main/Aoc2024/Day04/Solve.lean

Neat approach, I considered doing something like that briefly. Is List.diagonals in your Utils?

[James Sully]

Also I didn't know about #guard, that's super cool (docs)

[Matt Russell]

Neat approach, I considered doing something like that briefly. Is List.diagonals in your Utils?

Thanks! Yeah, I stole diagonals from a Haskell lib because I'm not smart enough to write the recursion myself: https://github.com/mdr/aoc-2024/blob/main/Aoc2024/Utils.lean#L89

[James Sully]

Ah yeah cool

[Kim Morrison]

Bonus points to make diagonals tail recursive. :-) (Even more points for the lemma relating it to the "easy to prove about" version.)

[Ayhon]

Kinda late to this one. This is my solution:

import Lean.Data.RBMap

namespace Day4

abbrev Input := Array (Array Char)

def parse_input(ls: List String): Input :=
  Array.mk $ ls
     |>.filter (not ∘ String.isEmpty)
     |>.map (Array.mk ·.trim.toList)

abbrev Board := Input

namespace Part1
  def List.iterate(f: α -> α)(a: α)(n: Nat) :=
    match n with
    | 0 => []
    | n'+1 => a :: List.iterate f (f a) n'

  def rot90[Inhabited α](b: Array (Array α)): Array (Array α) :=
    let n := b.size
    let m := b[0]!.size
    Array.ofFn fun (j: Fin m) =>
      Array.ofFn fun (i: Fin n) =>
        b[n-1-i]![j]!

  def count_dir(s: String)(dir: Bool × Bool)(b: Board): Nat :=
    let n := b.size - (if dir.1 then s.length-1 else 0)
    let m := b[0]!.size - (if dir.2 then s.length-1 else 0)
    List.ofFn (fun (i: Fin n) =>
      List.ofFn (fun (j: Fin m) =>
        if s.toList.enum.all (fun (o, c) =>
          let ni := if dir.1 then i + o else i
          let nj := if dir.2 then j + o else j
          b[ni]![nj]! == c
        )
        then 1
        else 0
      )
      |>.sum
    )
    |>.sum

  def solve(a: Input): Int :=
    let boards := List.iterate rot90 a 4
    let count_diag     := count_dir "XMAS" (true,  true)
    let count_straight := count_dir "XMAS" (false, true)
    boards
       |>.map (fun b => count_straight b + count_diag b)
       |>.sum
end Part1


namespace Part2
  abbrev Pattern := Array (Array (Option Char))

  def Pattern.count_match (p: Pattern)(b: Board) :=
    let bn := b.size
    let bm := b[0]!.size
    let pn := p.size
    let pm := p[0]!.size
    List.range (bn - pn + 1)
      |>.map (fun i =>
        List.range (bm - pm + 1)
          |>.map (fun j =>
            if ( List.range pn
              |>.all fun oi =>
                List.range pm
                  |>.all fun oj =>
                    match p[oi]![oj]! with
                    | none => true
                    | some c => c == b[i + oi]![j + oj]!)
            then 1 else 0
          )|>.sum
        )|>.sum


  def solve(a: Input): Int :=
    let pattern := #[
      #[some 'M', none,     some 'M' ],
      #[none,     some 'A', none ],
      #[some 'S', none,     some 'S' ]
    ]
    let patterns := Part1.List.iterate Part1.rot90 pattern 4
    patterns.map (fun p => Pattern.count_match p a) |>.sum

end Part2

[Kim Morrison]

Ouch rot90 scares me, with all that runtime array indexing. Can you use Vector?

[Kim Morrison]

Note for things like List.range pm |>.all you also have docs#Nat.all

[Ayhon]

Kim Morrison said:

Ouch rot90 scares me, with all that runtime array indexing. Can you use Vector?

Yes, that is fair. I have just spent too much time in a past project trying to convince Lean that array accesses were in fact in-bounds, and I didn't want to risk dealing with it in this Advent of Code.

I could use a Vector, but then I'd need to do checks in the parser that ensures the grid is well formed (all rows have the same length), and I'm not too sure how to achieve that. Any tips?

[James Sully]

I spent a fair bit of time banging my head against the wall of { grid : Array (Array Char) // rect grid } and variations on that theme.

[James Sully]

docs#Vector Oh cool, I had assumed Vectors were cons lists as in some other languages, not arrays

[Mario Carneiro]

Ayhon said:

I could use a Vector, but then I'd need to do checks in the parser that ensures the grid is well formed (all rows have the same length), and I'm not too sure how to achieve that. Any tips?

it's relatively easy to do such checks, you just use an if statement and fail if it's false. In the true case you get a hypothesis which you can adjoin to the array to build a Vector

[James Sully]

I don't find these things easy at all! Here's my attempt which I've gotten stuck on:

def String.lines (s : String) : List String := s.split (· == '\n')

-- not totally clear on whether the size should be implicit or not?
abbrev CharGrid {w h : Nat} := Vector (Vector Char w) h

def parse (input : String) : Option CharGrid := do
  let lines := input.lines.toArray
  if h₁ : lines.size > 0 then
    let width := lines[0].toList.toArray.size
    let mut ret := #[]
    for line in lines do
      let lineArr := line.toList.toArray
      if h₂ : lineArr.size = width then
        let lineVec : Vector Char width :=
          ⟨lineArr, by assumption⟩
        ret := ret.push lineVec
    have height := ret.size
    return ⟨ret, ???⟩
  else
    none

before I started trying to refactor it was

abbrev CharGrid := Array (Array Char)

def parse (input : String) : CharGrid :=
  let lines := input.lines.toArray
  lines.map (·.toList.toArray)

[Mario Carneiro]

What type is String.lines? That's not a MWE

[Mario Carneiro]

I guess it must be String -> List String

[James Sully]

sorry, I'll edit

[James Sully]

I'm having trouble connecting to the playground to verify it has everything that it needs, but I think it does, and in any case I've added String.lines

[Mario Carneiro]

def String.lines (s : String) : List String := s.split (· == '\n')

structure CharGrid where
  (w h : Nat)
  grid : Vector (Vector Char w) h

def parse (input : String) : Option CharGrid := do
  let first :: rest := input.lines | none
  let first := first.toList.toArray
  let w := first.size
  let rest : List (Vector _ w) ← rest.mapM fun s =>
    let s := s.toList.toArray
    if h : _ then some ⟨s, h⟩ else none
  let grid := (⟨first, rfl⟩ :: rest).toArray
  return { w, h := grid.size, grid := ⟨grid, rfl⟩ }

[James Sully]

Ahh, I've been wanting a way to do partial pattern matches

[James Sully]

ok, so you have to store the size at runtime

[James Sully]

thanks for this. To me it doesn't seem obviously worth the complication

[Mario Carneiro]

it depends on whether you want to have access to the width and height and/or validate that it's a rectangle

[Mario Carneiro]

not sure if that's part of the spec

[James Sully]

This line is sick
if h : _ then some ⟨s, h⟩ else none
I didn't know you could elide that. makes sense in retrospect

[Mario Carneiro]

you can even write

return { w, h := _, grid := ⟨grid, rfl⟩ }

on the last line

[Mario Carneiro]

you can write code where both data and control flow is clear only to lean itself :)

[Mario Carneiro]

heh

structure CharGrid where
  {w h : Nat}
  grid : Vector (Vector Char w) h

def parse (input : String) : Option CharGrid := do
  let first :: rest := input.lines | none
  let rest ← rest.mapM fun s =>
    if h : _ then some ⟨s.toList.toArray, h⟩ else none
  return ⟨⟨(⟨first.toList.toArray, rfl⟩ :: rest).toArray, rfl⟩⟩

[Ayhon]

The line I find the most magical is this one.

  let first :: rest := input.lines | none

If I understand correctly the semantics, the | none means that if the pattern is not matched, none is returned from the function directly. In some sense, it works similarly to the ? operator in Rust. Am I right? Also, is this syntax documented anywhere? Would like to know more (if there is more to know)

[Mario Carneiro]

this is equivalent to the let else syntax of rust

[Mario Carneiro]

let Cons(first, rest) = input.lines() else { return None };

[Mario Carneiro]

the ? operator in rust is more like the <- operator in lean

[Mario Carneiro]

In general, the syntax

do
  let pat := e
    | foo
  bar

is equivalent to

do
  match e with
  | pat => bar
  | _ => foo

[Greg Shuflin]

Mario Carneiro said:

you can even write

return { w, h := _, grid := ⟨grid, rfl⟩ }

on the last line

Where is this syntax with ⟨⟩ documented?

[Mario Carneiro]

not sure, but you might find references under the name "anonymous constructor (bracket)"

[Mario Carneiro]

I didn't find it in the WIP language reference, but there is a short syntax docstring if you hover it:

The anonymous constructor ⟨e, ...⟩ is equivalent to c e ... if the
expected type is an inductive type with a single constructor c.
If more terms are given than c has parameters, the remaining arguments
are turned into a new anonymous constructor application. For example,
⟨a, b, c⟩ : α × (β × γ) is equivalent to ⟨a, ⟨b, c⟩⟩.

[Paul Nelson]

abbrev Point := Nat × Nat
abbrev Direction := Int × Int
abbrev Grid := Array (Array Char)

def parse (input : String) : Grid :=
  input.splitOn "\n" |>.filter (· ≠ "") |>.map String.toList |>.map List.toArray |> List.toArray

def check (grid : Grid) (start : Point) (dir : Direction) (pattern : String) : Bool :=
  let (rows, cols) := (grid.size, grid[0]!.size)
  pattern.toList.enum.all fun ⟨i, t⟩ =>
    let (r, c) := (start.1 + i * dir.1, start.2 + i * dir.2)
    if r ≥ 0 && c ≥ 0 && r < rows && c < cols then
      let (r, c) := (r.toNat, c.toNat)
      grid[r]![c]! == t
    else
      false

def indices (grid : Grid) : List Point :=
  if grid.size = 0 then [] else
    let rows := List.range grid.size
    let cols := List.range grid[0]!.size
    rows.flatMap (cols.map fun j => ⟨·, j⟩)

def part1 (input : Grid) : Nat :=
  let dirs : List Direction := [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
  let pattern := "XMAS"
  indices input |>.flatMap (fun start =>
    dirs |>.filter (check input start · pattern))
    |>.length

def indicesWithRoom (grid : Grid) : List Point :=
  if grid.size = 0 then [] else
    let rows := List.range' 1 (grid.size - 2)
    let cols := List.range' 1 (grid[0]!.size - 2)
    rows.flatMap (cols.map fun j => ⟨·, j⟩)

def check2 (grid : Grid) (start : Point) : Bool :=
  let deltas : List Direction := [(1, 1), (-1, 1), (-1, -1), (1, -1)]
  let cornerChoices := ["MMSS", "MSSM", "SMMS", "SSMM"]
  let (r, c) := (start.1, start.2)
  grid[r]![c]! == 'A' && cornerChoices.any fun str =>
    deltas.enum.all fun ⟨i, delta⟩ =>
      let (r', c') := (r + delta.1, c + delta.2)
      let (r', c') := (r'.toNat, c'.toNat)
      grid[r']![c']! == str.toList[i]!

def part2 (input : Grid) : Nat :=
  indicesWithRoom input |>.filter (check2 input) |>.length