Advent of Code 2024 - Day 8

Stream: general

[Ayhon]

Posting my solution of Advent of code day 8, trying to encourage others to do the same. Didn't feel too functional today, so I let Id.run do most of the heavy lifting.

namespace Day8
open Std.Internal.Parsec.String
open Std.Internal.Parsec (skip many many1 attempt)

def andThen(a: α)(b: Prop)[Decidable b]: Option α := if b then some a else none

def List.allUnorderedPairs: List α -> List (α × α)
| x :: xs => xs.map ((x,·)) ++ List.allUnorderedPairs xs
| [] => []

abbrev Position := Nat × Nat
structure Input where
  antenas: Std.HashMap Char (Std.HashSet Position)
  bounds: Nat × Nat

def parse_input(ls: String): Input := Id.run do
  let mut antenas: Std.HashMap Char (Std.HashSet Position) := ∅
  let lines := ls.split (· == '\n') |>.filter (¬ ·.isEmpty)
  let positions := lines
    |>.enum
    |>.flatMap (fun (i,row) =>
      row.toList
        |>.enum
        |>.filterMap (fun (j, c) =>
          (c != '.') |> andThen (c,(i,j))
        )
    )
  for (c, (i,j)) in positions do
    antenas := antenas.getD c ∅
      |>.insert (i,j)
      |> Std.HashMap.insert antenas c
  return {antenas := antenas, bounds := (lines.length, lines[0]!.length)}

namespace Part1
  def solve(inp: Input): Int := Id.run do -- brrr
    let mut antinodes: Std.HashSet Position := ∅
    let getAntinodes (p1 p2: Position): List Position :=
      let (x1, y1) := (Int.ofNat p1.1, Int.ofNat p1.2)
      let (x2, y2) := (Int.ofNat p2.1, Int.ofNat p2.2)
      let (dx, dy) := (x2 - x1, y2 - y1)
      [(x2 + dx, y2 + dy), (x1 - dx, y1 - dy)]
        |>.filterMap (fun (x,y) =>
          (0 <= x && x < inp.bounds.1 && 0 <= y && y < inp.bounds.2)
            |> andThen ((x.toNat,y.toNat))
        )
    for (_, positions) in inp.antenas do
      for (p1,p2) in List.allUnorderedPairs positions.toList do
        antinodes := antinodes.insertMany $ getAntinodes p1 p2
    return antinodes.size
end Part1

namespace Part2
  def solve(inp: Input): Int := Id.run do -- brrr
    let mut antinodes: Std.HashSet Position := ∅
    let getAntinodes (p1 p2: Position): List Position :=
      let (x1, y1) := (Int.ofNat p1.1, Int.ofNat p1.2)
      let (x2, y2) := (Int.ofNat p2.1, Int.ofNat p2.2)
      let (dx, dy) := (x2 - x1, y2 - y1)
      let (w, h) := inp.bounds
      let rec getAntinodesInDirection (d: Int) (start: Int × Int)
      | 0 => panic "Too much"
      | lim+1 =>
        let (x,y) := (start.1 + d*dx, start.2 + d*dy)
        if 0 <= x && x < w && 0 <= y && y < h then
          (x.toNat,y.toNat) :: getAntinodesInDirection d (x,y) lim
        else []
      let limit := (w.max h)
      getAntinodesInDirection 1 (x1,y1) limit  ++ getAntinodesInDirection (-1) (x2,y2) limit
    for (_, positions) in inp.antenas do
      for (p1,p2) in List.allUnorderedPairs positions.toList do
        antinodes := antinodes.insertMany $ getAntinodes p1 p2
    return antinodes.size
end Part2

[Matt Russell]

https://github.com/mdr/aoc-2024/blob/main/Aoc2024/Day08/Solve.lean

[Paul Nelson]

import Lean.Data.HashSet

abbrev Grid := Array (Array Char)

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

abbrev Vec := Int × Int
instance : Add Vec where add p q := (p.1 + q.1, p.2 + q.2)
instance : Sub Vec where sub p q := (p.1 - q.1, p.2 - q.2)
instance : HMul Nat Vec Vec where hMul n p := (n * p.1, n * p.2)

def inBounds (height width : Nat) (p : Vec) : Bool :=
  let ⟨r, c⟩ := p; 0 ≤ r ∧ r < height ∧ 0 ≤ c ∧ c < width

def antinodes (height width : Nat) (a b : Vec) : List Vec :=
  let v := a - b; [a + v, b - v].filter (inBounds height width)

def Vec.primitiveMultiple (v : Vec) : Vec :=
  let d := Int.gcd v.1 v.2; (v.1 / d, v.2 / d)

def antinodes2 (height width : Nat) (a b : Vec) : List Vec :=
  let v := (a - b).primitiveMultiple
  let max_steps := max height width
  let forward := List.range max_steps
    |>.map (fun n => a + n * v)
    |>.takeWhile (inBounds height width)
  let backward := List.range' 1 max_steps
    |>.map (fun n => a - n * v)
    |>.takeWhile (inBounds height width)
  forward ++ backward

def Array.enum (ar : Array α) : Array (Nat × α) :=
  Array.mapIdx (fun i x => (i, x)) ar

def Grid.groupByChar (grid : Grid) : List (List Vec) :=
  let entries := grid.enum.flatMap fun (r, row) =>
    row.enum.filterMap fun (c, ch) =>
      if ch = '.' then none else some (ch, (r, c))
  entries.qsort (lt := (fun a b => a.1 < b.1))
    |>.toList.splitBy (·.1 == ·.1)
    |>.map (·.map (fun (_, (r,c)) => (Int.ofNat r, Int.ofNat c)))

def List.orderedPairs {α : Type} (xs : List α) : List (α × α) :=
  xs.enum.foldr (fun (i, x) acc =>
    xs.drop (i + 1) |>.map (fun y => (x, y)) |>.append acc
  ) []

def count (grid : Grid) (f : Nat → Nat → Vec → Vec → List Vec) : Nat :=
  grid.groupByChar
    |>.flatMap (·.orderedPairs)
    |>.flatMap (fun (a, b) => f grid.size grid[0]!.size a b)
    |>.foldl Std.HashSet.insert Std.HashSet.empty
    |>.size

def part1 (grid : Grid) : Nat := count grid antinodes

def part2 (grid : Grid) : Nat := count grid antinodes2