feat: Add enumeration between list indicies and element-multiplicity pairs.

Batteries/Data/List/Basic.lean

@@ -236,7 +236,8 @@ also receives each element's index added to an optional parameter `start`
 (i.e. the numbers that `f` takes as its first argument will be greater than or equal to `start` and
 less than `start + l.length`).
 -/
-@[specialize] def foldlIdx (f : Nat → α → β → α) (init : α) : List β → (start : Nat := 0) → α
+@[specialize] def foldlIdx {α : Type u} {β : Type v} (f : Nat → α → β → α) (init : α) :
+    List β → (start : Nat := 0) → α
   | [], _ => init
   | b :: l, s => foldlIdx f (f s init b) l (s + 1)
 
@@ -261,9 +262,9 @@ def foldrIdx {α : Type u} {β : Type v} (f : Nat → α → β → β) (init :
     (foldrIdx f i xs s, s) := by induction xs generalizing s <;> grind [foldrIdx]
   grind [foldrIdxTR]
 
-/-- `findIdxs p l` is the list of indexes of elements of `l` that satisfy `p`, added to an
-optional parameter `start` (so that the members of `findIdxs p l` will be greater than or
-equal to `start` and less than `l.length + start`).  -/
+/-- `findIdxs p l s` is the list of indexes of elements of `l` that satisfy `p`, added to an
+optional parameter `s` (so that the members of `findIdxs p l s` will be greater than or
+equal to `s` and less than `l.length + s`).  -/
 @[inline] def findIdxs (p : α → Bool) (l : List α) (start : Nat := 0) : List Nat :=
   foldrIdx (fun i a is => bif p a then i :: is else is) [] l start
 
@@ -278,9 +279,17 @@ We have `l.findIdxsValues p s = (l.findIdxs p s).zip (l.filter p)`.
 @[deprecated (since := "2025-11-06")]
 alias indexsValues := findIdxsValues
 
+/-- `findIdxNth p xs n` returns the index of the `n`th element for which `p` returns `true`. -/
+@[inline] def findIdxNth (p : α → Bool) (xs : List α) (n : Nat) : Nat := go xs n 0 where
+  /-- Auxiliary for `findIdxNth`: `findIdxNth.go p l n acc = findIdxNth p l n + acc`. -/
+  @[specialize] go : (xs : List α) → (n : Nat) → (s : Nat) → Nat
+  | [], _, s => s
+  | a :: xs, 0, s => bif p a then s else go xs 0 (s + 1)
+  | a :: xs, n + 1, s => go xs (n + bif !(p a) then 1 else 0) (s + 1)
+
 /--
-`idxsOf a l` is the list of all indexes of `a` in `l`,  added to an
-optional parameter `start`. For example:
+`idxsOf a l s` is the list of all indexes of `a` in `l`,  added to an
+optional parameter `s`. For example:
 ```
 idxsOf b [a, b, a, a] = [1]
 idxsOf a [a, b, a, a] 5 = [5, 7, 8]
@@ -292,6 +301,24 @@ idxsOf a [a, b, a, a] 5 = [5, 7, 8]
 @[deprecated (since := "2025-11-06")]
 alias indexesOf := idxsOf
 
+/-- `idxOfNth p xs n` returns the index of the `n`th instance of `a` in `xs`. -/
+def idxOfNth [BEq α] (a : α) (xs : List α) (n : Nat) : Nat :=
+  xs.findIdxNth (· == a) n
+
+/-- `countPBefore p xs i hip` counts the number of `x` in `xs` before the `i`th index for
+which `p x = true`. -/
+def countPBefore (p : α → Bool) (xs : List α) (i : Nat) : Nat := go xs i 0 where
+  /-- Auxiliary for `countPBefore`: `countPBefore.go p l i acc = countPBefore p l i + acc`. -/
+  @[specialize] go : (xs : List α) → (i : Nat) → (s : Nat) → Nat
+  | _ :: _, 0, s => s
+  | a :: xs, i + 1, s => go xs i (bif p a then s + 1 else s)
+  | [], _, s => s
+
+/-- `countBefore x xs n` counts the number of `x` in `xs` before the
+    `i`th index for which `x == a` is true. -/
+def countBefore [BEq α] (a : α) : List α → Nat → Nat :=
+  countPBefore (· == a)
+
 /--
 `lookmap` is a combination of `lookup` and `filterMap`.
 `lookmap f l` will apply `f : α → Option α` to each element of the list,

Batteries/Data/List/Count.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Parikshit Khanna, Jeremy Avigad, Leonardo de Moura, Floris van Doorn, Mario Carneiro
 -/
 module
+public import Batteries.Data.List.Lemmas
 
 @[expose] public section
 
@@ -20,13 +21,51 @@ open Nat
 
 namespace List
 
-
 /-! ### count -/
 
 section count
 
-variable [DecidableEq α]
+theorem count_singleton' [DecidableEq α] (a b : α) : count a [b] = if b = a then 1 else 0 :=
+  count_singleton.trans (by grind)
+
+theorem count_concat [BEq α] [LawfulBEq α] (a : α) (l : List α) :
+    count a (concat l a) = succ (count a l) := by simp
+
+/-! ### idxToSigmaCount, sigmaCountToIdx -/
+
+/-- `idxToSigmaCount` is a `Fin`-to-`Fin` wrapper for `countBefore`. -/
+def idxToSigmaCount [BEq α] [ReflBEq α] (xs : List α) (i : Fin (xs.length)) :
+    (x : α) × Fin (xs.count x) :=
+  ⟨xs[i.1], xs.countBefore xs[i.1] i, countBefore_lt_count_of_lt_length_of_beq _ BEq.rfl⟩
+
+@[simp, grind =]
+theorem fst_idxToSigmaCount [BEq α] [ReflBEq α] {xs: List α} {i : Fin (xs.length)} :
+    (xs.idxToSigmaCount i).1 = xs[i.1] := rfl
+
+@[simp, grind =]
+theorem snd_idxToSigmaCount [BEq α] [ReflBEq α] {xs: List α} {i : Fin (xs.length)} :
+    (xs.idxToSigmaCount i).2 =
+    ⟨xs.countBefore xs[i.1] i, countBefore_lt_count_of_lt_length_of_beq _ BEq.rfl⟩ := rfl
+
+
+@[simp, grind =]
+theorem coe_snd_idxToSigmaCount [BEq α] [ReflBEq α] {xs: List α} {i : Fin (xs.length)} :
+    ((xs.idxToSigmaCount i).2 : Nat) = xs.countBefore xs[i.1] i := rfl
+
+/-- `sigmaCountToIdx` is a `_ × Fin`-to-`Fin` wrapper for `countBefore`. -/
+def sigmaCountToIdx [BEq α] (xs: List α) (xc : (x : α) × Fin (xs.count x)) :
+    Fin (xs.length) := ⟨xs.idxOfNth xc.1 xc.2, idxOfNth_lt_length_of_lt_count xc.2.isLt⟩
+
+@[simp, grind =]
+theorem coe_sigmaCountToIdx [BEq α] {xs: List α} {xc : (x : α) × Fin (xs.count x)} :
+    (xs.sigmaCountToIdx xc : Nat) = xs.idxOfNth xc.1 xc.2 := rfl
 
-theorem count_singleton' (a b : α) : count a [b] = if b = a then 1 else 0 := by simp [count_cons]
+@[simp, grind =]
+theorem sigmaCountToIdx_idxToSigmaCount [BEq α] [ReflBEq α] {xs : List α}
+    {i : Fin (xs.length)} : xs.sigmaCountToIdx (xs.idxToSigmaCount i) = i :=
+  Fin.ext <| idxOfNth_countBefore_of_lt_length_of_beq _ BEq.rfl
 
-theorem count_concat (a : α) (l : List α) : count a (concat l a) = succ (count a l) := by simp
+@[simp, grind =]
+theorem idxToSigmaCount_sigmaCountToIdx [BEq α] [LawfulBEq α] {xs : List α}
+    {xc : (x : α) × Fin (xs.count x)} : xs.idxToSigmaCount (xs.sigmaCountToIdx xc) = xc :=
+  Sigma.ext getElem_idxOfNth_eq (heq_of_eqRec_eq (by simp) (by grind))