From e6f4111aa97f34383369577bff011d546d26d000 Mon Sep 17 00:00:00 2001
From: Andreas Grois <andi@grois.info>
Date: Thu, 7 Dec 2023 21:53:19 +0100
Subject: Day 6

---
 Common/Option.lean |  4 +++
 Day6.lean          | 84 ++++++++++++++++++++++++++++++++++++++++++++++++------
 Main.lean          |  3 ++
 3 files changed, 82 insertions(+), 9 deletions(-)

diff --git a/Common/Option.lean b/Common/Option.lean
index 23ea8d6..275f83f 100644
--- a/Common/Option.lean
+++ b/Common/Option.lean
@@ -5,3 +5,7 @@ def zip (a : Option α) (b : Option β) : Option (α×β) := a >>= λ a ↦ b >>
 def unzip : (a : Option (α×β)) → (Option α) × (Option β)
 | none => (none, none)
 | some (a, b) => (some a, some b)
+
+def toExcept {α : Type u0} {ε : Type u1} { m : Type u0 → Type u2} [Pure m] [MonadExcept m (ε := ε)] (error : ε) : Option α → m α
+| none => throw error
+| some a => pure a
diff --git a/Day6.lean b/Day6.lean
index 54464ea..93f7dbb 100644
--- a/Day6.lean
+++ b/Day6.lean
@@ -1,27 +1,93 @@
 import Common.Option
 import Common.Helpers
+import Common.DayPart
 
 structure Race where
   timeLimit : Nat
   recordDistance : Nat
+  deriving Repr
 
-private def parseLine (header : String) (input : String) : Option (List Nat) := do
+private def parseLine (header : String) (input : String) : Except String (List Nat) := do
   if not $ input.startsWith header then
-    failure
+    throw s!"Unexpected line header: {header}, {input}"
   let input := input.drop header.length |> String.trim
   let numbers := input.split Char.isWhitespace
     |> List.map String.trim
-    |> List.filter String.isEmpty
-  numbers.mapM String.toNat?
+    |> List.filter (not ∘ String.isEmpty)
+  numbers.mapM $ Option.toExcept s!"Failed to parse input line: Not a number {input}" ∘  String.toNat?
 
-def parse (input : String) : Option (List Race) := do
+def parse (input : String) : Except String (List Race) := do
   let lines := input.splitOn "\n"
     |> List.map String.trim
-    |> List.filter String.isEmpty
+    |> List.filter (not ∘ String.isEmpty)
   let (times, distances) ← match lines with
-    | [times, distances] => Option.zip (parseLine "Time:" times) (parseLine "Distance:" distances)
-    | _ => none
+    | [times, distances] =>
+      let times ← parseLine "Time:" times
+      let distances ← parseLine "Distance:" distances
+      pure (times, distances)
+    | _ => throw "Failed to parse: there should be exactly 2 lines of input"
   if times.length != distances.length then
-    failure
+    throw "Input lines need to have the same number of, well, numbers."
   let pairs := times.zip distances
+  if pairs = [] then
+    throw "Input does not have at least one race."
   return pairs.map $ uncurry Race.mk
+
+-- okay, part 1 is a quadratic equation. Simple as can be
+-- s = v * tMoving
+-- s = tPressed * (tLimit - tPressed)
+-- (tPressed - tLimit) * tPressed + s = 0
+-- tPressed² - tPressed * tLimit + s = 0
+-- tPressed := tLimit / 2 ± √(tLimit² / 4 - s)
+-- beware: We need to _beat_ the record, so s here is the record + 1
+
+-- Inclusive! This is the smallest number that can win, and the largest number that can win
+private def Race.timeRangeToWin (input : Race) : (Nat × Nat) :=
+  let tLimit  := input.timeLimit.toFloat
+  let sRecord := input.recordDistance.toFloat
+  let tlimitHalf := 0.5 * tLimit
+  let theRoot := (tlimitHalf^2 - sRecord - 1.0).sqrt
+  let lowerBound := tlimitHalf - theRoot
+  let upperBound := tlimitHalf + theRoot
+  let lowerBound := lowerBound.ceil.toUInt64.toNat
+  let upperBound := upperBound.floor.toUInt64.toNat
+  (lowerBound,upperBound)
+
+def part1 (input : List Race) : Nat :=
+  let limits := input.map Race.timeRangeToWin
+  let counts := limits.map $ λ p ↦ p.snd - p.fst + 1 -- inclusive range
+  counts.foldl (· * ·) 1
+
+-- part2 is the same thing, but here we need to be careful.
+-- namely, careful about the precision of Float. Which luckily is enough, as confirmed by pen&paper
+-- but _barely_ enough.
+-- If Lean's Float were an actual C float and not a C double, this would not work.
+
+-- we need to concatenate the numbers again (because I don't want to make a separate parse for part2)
+private def combineNumbers (left : Nat) (right : Nat) : Nat :=
+  let rec countDigits := λ (s : Nat) (n : Nat) ↦
+    if p : n > 0 then
+      have : n > n / 10 := by simp[p, Nat.div_lt_self]
+      countDigits (s+1) (n/10)
+    else
+      s
+  let d := if right = 0 then 1 else countDigits 0 right
+  left * (10^d) + right
+
+def part2 (input : List Race) : Nat :=
+  let timeLimits := input.map Race.timeLimit
+  let timeLimit := timeLimits.foldl combineNumbers 0
+  let records := input.map Race.recordDistance
+  let record := records.foldl combineNumbers 0
+  let limits := Race.timeRangeToWin $ {timeLimit := timeLimit, recordDistance := record}
+  limits.snd - limits.fst + 1 -- inclusive range
+
+open DayPart
+instance : Parse ⟨6, by simp⟩ (ι := List Race) where
+  parse := parse
+
+instance : Part ⟨6, _⟩ Parts.One (ι := List Race) (ρ := Nat) where
+  run := some ∘ part1
+
+instance : Part ⟨6, _⟩ Parts.Two (ι := List Race) (ρ := Nat) where
+  run := some ∘ part2
diff --git a/Main.lean b/Main.lean
index 1ad3933..7bfce37 100644
--- a/Main.lean
+++ b/Main.lean
@@ -4,6 +4,7 @@ import «Day2»
 import «Day3»
 import «Day4»
 import «Day5»
+import «Day6»
 
 open DayPart
 
@@ -26,6 +27,8 @@ def try_run_day_part (day : Days) (part : Parts) (data : String) : IO String :=
   | ⟨4,_⟩, Parts.Two => impl ⟨4,_⟩ Parts.Two data
   | ⟨5,_⟩, Parts.One => impl ⟨5,_⟩ Parts.One data
   | ⟨5,_⟩, Parts.Two => impl ⟨5,_⟩ Parts.Two data
+  | ⟨6,_⟩, Parts.One => impl ⟨6,_⟩ Parts.One data
+  | ⟨6,_⟩, Parts.Two => impl ⟨6,_⟩ Parts.Two data
   | _, _ => throw $ IO.userError "The requested combination of day/part has not been implemented yet."
 
 def main (parameters : List String): IO Unit := do
-- 
cgit v1.2.3