Advent of Code 2024 - Day 7

Stream: general

[Ayhon]

This is my solution for day 7. In general I'm much more satisfied with this solution than the one from the previous day.

import Std.Internal.Parsec
import Std.Internal.Parsec.String
import Std.Data.HashSet

namespace Day7
open Std.Internal.Parsec.String (pchar pstring skipString skipChar Parser digits ws)
open Std.Internal.Parsec (skip many many1 attempt)

inductive Operation where
| Sum
| Mul
| Concat
deriving BEq, Repr, Inhabited

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

def Nat.isSuffix (big small: Nat): Option Nat := loop big big small where
  loop : Nat -> Nat -> Nat -> Option Nat
  | 0, _, _ => none
  | limit+1, a, b =>
    if (a % 10 == b % 10) then (((b < 10) |> andThen (a/10)) <|> (loop limit (a/10) (b/10))) else none

def Operation.undo(op: Operation)(result operand: Nat): Option Nat := do
  match op with
  | Sum =>
    (result >= operand) |> andThen (result - operand)
  | Mul =>
    (result % operand == 0) |> andThen (result / operand)
  | Concat => Nat.isSuffix result operand

structure Equation where
  result: Nat
  operands: Array Nat
deriving BEq, Repr, Inhabited

def Equation.parser: Parser Equation := do
  let result <- digits
  skipString ": "
  let operands <- many (digits <* (attempt (skipChar ' ') <|> pure ()))
  return {result, operands: Equation}

abbrev Input := Array Equation

def parse_input(content: String): Option Input :=
  let parser := many (Equation.parser <* skipChar '\n')
  match parser content.iter with
  | .success _ equations => some equations
  | .error _ reason => dbg_trace reason; none

def HashSet.flatMap [Hashable α][BEq α][Hashable β][BEq β](f: α -> List β): Std.HashSet α -> Std.HashSet β :=
  Std.HashSet.fold (fun (newSet: Std.HashSet β) (elem: α) =>
    newSet.insertMany $ f elem
  ) ∅

def Equation.isValid(operations: List Operation)(eq: Equation): Bool := go eq.operands.reverse.toList {eq.result} where
  go : List Nat -> Std.HashSet Nat -> Bool
  | [], _ => false
  | [x], possibleResults => possibleResults.contains x
  | x :: xs, possibleResults =>
    let getNewPossibleResults result := operations.map (Operation.undo · result x) |>.reduceOption
    go xs $ HashSet.flatMap getNewPossibleResults possibleResults

namespace Part1
  def solve(inp: Option Input): Int := Id.run do
    let some equations := inp | pure $ -1 -- Can't use -1 on its own
    return equations
      |>.filter (Equation.isValid [Operation.Sum, Operation.Mul])
      |>.map (·.result)
      |>.toList
      |>.sum
end Part1

namespace Part2
  def solve(inp: Option Input): Int := Id.run do
    let some equations := inp | pure $ -1 -- Can't use -1 on its own
    return equations
      |>.filter (Equation.isValid [Operation.Sum, Operation.Mul, Operation.Concat])
      |>.map (·.result)
      |>.toList
      |>.sum
end Part2

[Ayhon]

I'll try to find some time to work on day 8 tomorrow

[Patrick Massot]

I only took a very quick look and I don’t know what the problem is about, but the pattern |>.filter … |>.map … is a bit scary. This is only a linear cost of course, but traversing the array twice feels a bit wasteful. There is some docs#Array.filterMap whose name is promising. Of course one could argue that you should also do the sum in the same loop, and now there is probably no builtin function so you would need to think a bit more about the function to plug into filterMap.

[Ayhon]

Patrick Massot said:

I only took a very quick look and I don’t know what the problem is about, but the pattern |>.filter … |>.map … is a bit scary. This is only a linear cost of course, but traversing the array twice feels a bit wasteful. There is some docs#Array.filterMap whose name is promising. Of course one could argue that you should also do the sum in the same loop, and now there is probably no builtin function so you would need to think a bit more about the function to plug into filterMap.

I'm actualy not too concerned about efficiency in this case. If this were a property I cared about though, I guess I could implement all of this inside a fold, with a function that does all of these operations at the same time. However, I believe that the sacrifice in readability is not worth the cost. Furthermore, I believe that a sufficiently good compiler should merge all of these operations together, since these compositions form some sort of pipeline.

[Patrick Massot]

That’s a good question. I have no idea whether Lean’s compiler does such things. My guess would be that the current compiler does not but the new one will. @Henrik Böving, any insight?

[Henrik Böving]

This optimization is known as fusion in Haskell land and indeed not implemented by the current compiler. I don't know about the plans we are going for with the new compiler with respect to this. Note that as Rust's iterator interface demonstrated you do not necessarily need to teach your compiler about fusion in order to get fusion if your iteration primitives are sufficiently well designed (of course the compiler also plays somewhat of a role here but it is not as crazy as in Haskell)

[Ayhon]

Admittedly, when I was talking about the compiler being able to do these kinds of optimizations, I was anecdotically quoting what I recalled haskell did with fusion. I know it did some things with map and other primitive operations, but I wasn't completely certain that it worked with filter and map in this way. Good to know I wasn't far off the mark.

[Ayhon]

Also, @Patrick Massot , after thinking a bit more about this, I've actually identified some instances where I used reduceOption but a filterMap could have been more idiomatic. Thank you for bringing it up to my attention, and apologies for not thanking you with my first response

[Patrick Massot]

No problem. My answer was not meant to be deep anyway, only a comment after a very superficial reading in case you weren’t aware of filterMap.

[Paul Nelson]

structure Problem where
  target: Nat
  inputs: List Nat
  deriving Repr

def parse (input : String) : List Problem :=
  let lines := input.splitOn "\n" |>.filter (!·.isEmpty)
  lines.map fun line =>
    let parts := line.splitOn ": "
    let target := parts[0]!.toNat!
    let inputs := parts[1]!.splitOn " " |>.map String.toNat!
    { target := target, inputs := inputs }

def Nat.subtract? (y x : Nat) : Option Nat :=
  if y >= x then some (y - x) else none

def Nat.divide? (y x : Nat) : Option Nat :=
  if y % x == 0 then some (y / x) else none

def solvable (inv_ops : List (Nat → Nat → Option Nat)) (problem : Problem) : Bool :=
  let rec go (rev_inputs : List Nat) (y : Nat) : Bool :=
    match rev_inputs with
    | [] => false
    | [x] => x == y
    | x::xs => inv_ops.any fun op =>
        op y x |>.map (go xs) |>.getD false
  go problem.inputs.reverse problem.target

def part1 (problems : List Problem) : Nat :=
  let ops := [Nat.subtract?, Nat.divide?]
  problems |>.filter (solvable ops) |>.map (·.target) |>.sum

def Nat.dropSuffix? (y x: Nat) : Option Nat := do
  let y'_str ← (y.repr).dropSuffix? x.repr
  pure (if y'_str.isEmpty then 0 else y'_str.toString.toNat!)

def part2 (problems : List Problem) : Nat :=
  let ops := [Nat.subtract?, Nat.divide?, Nat.dropSuffix?]
  problems |>.filter (solvable ops) |>.map (·.target) |>.sum