From ec7ae594731964874d6b29f06ba557d60ca7e3b5 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Mon, 9 Sep 2024 17:04:57 +1000
Subject: [PATCH] feat: List.count lemmas (#5285)

---
 src/Init/Data/List/Count.lean     | 98 ++++++++++++++++++++++++++++---
 src/Init/Data/List/Nat.lean       |  1 +
 src/Init/Data/List/Nat/Basic.lean |  2 +-
 src/Init/Data/List/Nat/Count.lean | 31 ++++++++++
 src/Init/Data/List/Perm.lean      |  2 +-
 src/Init/Data/Option/Lemmas.lean  | 26 ++++++++
 6 files changed, 151 insertions(+), 9 deletions(-)
 create mode 100644 src/Init/Data/List/Nat/Count.lean

diff --git a/src/Init/Data/List/Count.lean b/src/Init/Data/List/Count.lean
index c39390b88517..41e57b44b9fa 100644
--- a/src/Init/Data/List/Count.lean
+++ b/src/Init/Data/List/Count.lean
@@ -40,6 +40,9 @@ protected theorem countP_go_eq_add (l) : countP.go p l n = n + countP.go p l 0 :
 theorem countP_cons (a : α) (l) : countP p (a :: l) = countP p l + if p a then 1 else 0 := by
   by_cases h : p a <;> simp [h]
 
+theorem countP_singleton (a : α) : countP p [a] = if p a then 1 else 0 := by
+  simp [countP_cons]
+
 theorem length_eq_countP_add_countP (l) : length l = countP p l + countP (fun a => ¬p a) l := by
   induction l with
   | nil => rfl
@@ -61,6 +64,10 @@ theorem countP_eq_length_filter (l) : countP p l = length (filter p l) := by
     then rw [countP_cons_of_pos p l h, ih, filter_cons_of_pos h, length]
     else rw [countP_cons_of_neg p l h, ih, filter_cons_of_neg h]
 
+theorem countP_eq_length_filter' : countP p = length ∘ filter p := by
+  funext l
+  apply countP_eq_length_filter
+
 theorem countP_le_length : countP p l ≤ l.length := by
   simp only [countP_eq_length_filter]
   apply length_filter_le
@@ -68,15 +75,38 @@ theorem countP_le_length : countP p l ≤ l.length := by
 @[simp] theorem countP_append (l₁ l₂) : countP p (l₁ ++ l₂) = countP p l₁ + countP p l₂ := by
   simp only [countP_eq_length_filter, filter_append, length_append]
 
-theorem countP_pos {p} : 0 < countP p l ↔ ∃ a ∈ l, p a := by
+@[simp] theorem countP_pos_iff {p} : 0 < countP p l ↔ ∃ a ∈ l, p a := by
   simp only [countP_eq_length_filter, length_pos_iff_exists_mem, mem_filter, exists_prop]
 
-theorem countP_eq_zero {p} : countP p l = 0 ↔ ∀ a ∈ l, ¬p a := by
+@[deprecated countP_pos_iff (since := "2024-09-09")] abbrev countP_pos := @countP_pos_iff
+
+@[simp] theorem one_le_countP_iff {p} : 1 ≤ countP p l ↔ ∃ a ∈ l, p a :=
+  countP_pos_iff
+
+@[simp] theorem countP_eq_zero {p} : countP p l = 0 ↔ ∀ a ∈ l, ¬p a := by
   simp only [countP_eq_length_filter, length_eq_zero, filter_eq_nil_iff]
 
-theorem countP_eq_length {p} : countP p l = l.length ↔ ∀ a ∈ l, p a := by
+@[simp] theorem countP_eq_length {p} : countP p l = l.length ↔ ∀ a ∈ l, p a := by
   rw [countP_eq_length_filter, filter_length_eq_length]
 
+theorem countP_replicate (p : α → Bool) (a : α) (n : Nat) :
+    countP p (replicate n a) = if p a then n else 0 := by
+  simp only [countP_eq_length_filter, filter_replicate]
+  split <;> simp
+
+theorem boole_getElem_le_countP (p : α → Bool) (l : List α) (i : Nat) (h : i < l.length) :
+    (if p l[i] then 1 else 0) ≤ l.countP p := by
+  induction l generalizing i with
+  | nil => simp at h
+  | cons x l ih =>
+    cases i with
+    | zero => simp [countP_cons]
+    | succ i =>
+      simp only [length_cons, add_one_lt_add_one_iff] at h
+      simp only [getElem_cons_succ, countP_cons]
+      specialize ih _ h
+      exact le_add_right_of_le ih
+
 theorem Sublist.countP_le (s : l₁ <+ l₂) : countP p l₁ ≤ countP p l₂ := by
   simp only [countP_eq_length_filter]
   apply s.filter _ |>.length_le
@@ -102,6 +132,30 @@ theorem countP_filter (l : List α) :
   | [] => rfl
   | a :: l => by rw [map_cons, countP_cons, countP_cons, countP_map p f l]; rfl
 
+theorem length_filterMap_eq_countP (f : α → Option β) (l : List α) :
+    (filterMap f l).length = countP (fun a => (f a).isSome) l := by
+  induction l with
+  | nil => rfl
+  | cons x l ih =>
+    simp only [filterMap_cons, countP_cons]
+    split <;> simp [ih, *]
+
+theorem countP_filterMap (p : β → Bool) (f : α → Option β) (l : List α) :
+    countP p (filterMap f l) = countP (fun a => ((f a).map p).getD false) l := by
+  simp only [countP_eq_length_filter, filter_filterMap, ← filterMap_eq_filter]
+  simp only [length_filterMap_eq_countP]
+  congr
+  ext a
+  simp (config := { contextual := true }) [Option.getD_eq_iff]
+
+@[simp] theorem countP_join (l : List (List α)) :
+    countP p l.join = Nat.sum (l.map (countP p)) := by
+  simp only [countP_eq_length_filter, filter_join]
+  simp [countP_eq_length_filter']
+
+@[simp] theorem countP_reverse (l : List α) : countP p l.reverse = countP p l := by
+  simp [countP_eq_length_filter, filter_reverse]
+
 variable {p q}
 
 theorem countP_mono_left (h : ∀ x ∈ l, p x → q x) : countP p l ≤ countP q l := by
@@ -136,6 +190,11 @@ theorem count_cons (a b : α) (l : List α) :
     count a (b :: l) = count a l + if b == a then 1 else 0 := by
   simp [count, countP_cons]
 
+theorem count_eq_countP (a : α) (l : List α) : count a l = countP (· == a) l := rfl
+theorem count_eq_countP' {a : α} : count a = countP (· == a) := by
+  funext l
+  apply count_eq_countP
+
 theorem count_tail : ∀ (l : List α) (a : α) (h : l ≠ []),
       l.tail.count a = l.count a - if l.head h == a then 1 else 0
   | head :: tail, a, _ => by simp [count_cons]
@@ -157,6 +216,17 @@ theorem count_singleton (a b : α) : count a [b] = if b == a then 1 else 0 := by
 @[simp] theorem count_append (a : α) : ∀ l₁ l₂, count a (l₁ ++ l₂) = count a l₁ + count a l₂ :=
   countP_append _
 
+theorem count_join (a : α) (l : List (List α)) : count a l.join = Nat.sum (l.map (count a)) := by
+  simp only [count_eq_countP, countP_join, count_eq_countP']
+
+@[simp] theorem count_reverse (a : α) (l : List α) : count a l.reverse = count a l := by
+  simp only [count_eq_countP, countP_eq_length_filter, filter_reverse, length_reverse]
+
+theorem boole_getElem_le_count (a : α) (l : List α) (i : Nat) (h : i < l.length) :
+    (if l[i] == a then 1 else 0) ≤ l.count a := by
+  rw [count_eq_countP]
+  apply boole_getElem_le_countP (· == a)
+
 variable [LawfulBEq α]
 
 @[simp] theorem count_cons_self (a : α) (l : List α) : count a (a :: l) = count a l + 1 := by
@@ -172,14 +242,19 @@ theorem count_concat_self (a : α) (l : List α) :
     count a (concat l a) = (count a l) + 1 := by simp
 
 @[simp]
-theorem count_pos_iff_mem {a : α} {l : List α} : 0 < count a l ↔ a ∈ l := by
-  simp only [count, countP_pos, beq_iff_eq, exists_eq_right]
+theorem count_pos_iff {a : α} {l : List α} : 0 < count a l ↔ a ∈ l := by
+  simp only [count, countP_pos_iff, beq_iff_eq, exists_eq_right]
+
+@[deprecated count_pos_iff (since := "2024-09-09")] abbrev count_pos_iff_mem := @count_pos_iff
+
+@[simp] theorem one_le_count_iff {a : α} {l : List α} : 1 ≤ count a l ↔ a ∈ l :=
+  count_pos_iff
 
 theorem count_eq_zero_of_not_mem {a : α} {l : List α} (h : a ∉ l) : count a l = 0 :=
-  Decidable.byContradiction fun h' => h <| count_pos_iff_mem.1 (Nat.pos_of_ne_zero h')
+  Decidable.byContradiction fun h' => h <| count_pos_iff.1 (Nat.pos_of_ne_zero h')
 
 theorem not_mem_of_count_eq_zero {a : α} {l : List α} (h : count a l = 0) : a ∉ l :=
-  fun h' => Nat.ne_of_lt (count_pos_iff_mem.2 h') h.symm
+  fun h' => Nat.ne_of_lt (count_pos_iff.2 h') h.symm
 
 theorem count_eq_zero {l : List α} : count a l = 0 ↔ a ∉ l :=
   ⟨not_mem_of_count_eq_zero, count_eq_zero_of_not_mem⟩
@@ -224,6 +299,15 @@ theorem count_le_count_map [DecidableEq β] (l : List α) (f : α → β) (x : 
   rw [count, count, countP_map]
   apply countP_mono_left; simp (config := { contextual := true })
 
+theorem count_filterMap {α} [BEq β] (b : β) (f : α → Option β) (l : List α) :
+    count b (filterMap f l) = countP (fun a => f a == some b) l := by
+  rw [count_eq_countP, countP_filterMap]
+  congr
+  ext a
+  obtain _ | b := f a
+  · simp
+  · simp
+
 theorem count_erase (a b : α) :
     ∀ l : List α, count a (l.erase b) = count a l - if b == a then 1 else 0
   | [] => by simp
diff --git a/src/Init/Data/List/Nat.lean b/src/Init/Data/List/Nat.lean
index 877061671f34..b13a173cb21e 100644
--- a/src/Init/Data/List/Nat.lean
+++ b/src/Init/Data/List/Nat.lean
@@ -9,3 +9,4 @@ import Init.Data.List.Nat.Pairwise
 import Init.Data.List.Nat.Range
 import Init.Data.List.Nat.Sublist
 import Init.Data.List.Nat.TakeDrop
+import Init.Data.List.Nat.Count
diff --git a/src/Init/Data/List/Nat/Basic.lean b/src/Init/Data/List/Nat/Basic.lean
index 77cee69a9ee2..4e2f97faf6e9 100644
--- a/src/Init/Data/List/Nat/Basic.lean
+++ b/src/Init/Data/List/Nat/Basic.lean
@@ -23,7 +23,7 @@ namespace List
 theorem length_filter_lt_length_iff_exists {l} :
     length (filter p l) < length l ↔ ∃ x ∈ l, ¬p x := by
   simpa [length_eq_countP_add_countP p l, countP_eq_length_filter] using
-    countP_pos (p := fun x => ¬p x)
+    countP_pos_iff (p := fun x => ¬p x)
 
 /-! ### reverse -/
 
diff --git a/src/Init/Data/List/Nat/Count.lean b/src/Init/Data/List/Nat/Count.lean
new file mode 100644
index 000000000000..4ac93bdce29b
--- /dev/null
+++ b/src/Init/Data/List/Nat/Count.lean
@@ -0,0 +1,31 @@
+/-
+Copyright (c) 2024 Kim Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+prelude
+import Init.Data.List.Count
+import Init.Data.Nat.Lemmas
+
+namespace List
+
+open Nat
+
+theorem countP_set (p : α → Bool) (l : List α) (i : Nat) (a : α) (h : i < l.length) :
+    (l.set i a).countP p = l.countP p - (if p l[i] then 1 else 0) + (if p a then 1 else 0) := by
+  induction l generalizing i with
+  | nil => simp at h
+  | cons x l ih =>
+    cases i with
+    | zero => simp [countP_cons]
+    | succ i =>
+      simp [add_one_lt_add_one_iff] at h
+      simp [countP_cons, ih _ h]
+      have : (if p l[i] = true then 1 else 0) ≤ l.countP p := boole_getElem_le_countP p l i h
+      omega
+
+theorem count_set [BEq α] (a b : α) (l : List α) (i : Nat) (h : i < l.length) :
+    (l.set i a).count b = l.count b - (if l[i] == b then 1 else 0) + (if a == b then 1 else 0) := by
+  simp [count_eq_countP, countP_set, h]
+
+end List
diff --git a/src/Init/Data/List/Perm.lean b/src/Init/Data/List/Perm.lean
index e491572c9c65..ec31e51319ab 100644
--- a/src/Init/Data/List/Perm.lean
+++ b/src/Init/Data/List/Perm.lean
@@ -348,7 +348,7 @@ theorem perm_iff_count {l₁ l₂ : List α} : l₁ ~ l₂ ↔ ∀ a, count a l
       specialize H b
       simp at H
   | cons a l₁ IH =>
-    have : a ∈ l₂ := count_pos_iff_mem.mp (by rw [← H]; simp)
+    have : a ∈ l₂ := count_pos_iff.mp (by rw [← H]; simp)
     refine ((IH fun b => ?_).cons a).trans (perm_cons_erase this).symm
     specialize H b
     rw [(perm_cons_erase this).count_eq] at H
diff --git a/src/Init/Data/Option/Lemmas.lean b/src/Init/Data/Option/Lemmas.lean
index ed5af61eb234..345086d9bc3f 100644
--- a/src/Init/Data/Option/Lemmas.lean
+++ b/src/Init/Data/Option/Lemmas.lean
@@ -78,6 +78,9 @@ theorem eq_none_iff_forall_not_mem : o = none ↔ ∀ a, a ∉ o :=
 
 theorem isSome_iff_exists : isSome x ↔ ∃ a, x = some a := by cases x <;> simp [isSome]
 
+@[simp] theorem isSome_eq_isSome : (isSome x = isSome y) ↔ (x = none ↔ y = none) := by
+  cases x <;> cases y <;> simp
+
 @[simp] theorem isNone_none : @isNone α none = true := rfl
 
 @[simp] theorem isNone_some : isNone (some a) = false := rfl
@@ -220,8 +223,17 @@ theorem mem_map_of_mem (g : α → β) (h : a ∈ x) : g a ∈ Option.map g x :=
   split <;> rfl
 
 @[simp] theorem filter_none (p : α → Bool) : none.filter p = none := rfl
+
 theorem filter_some : Option.filter p (some a) = if p a then some a else none := rfl
 
+theorem isSome_filter_of_isSome (p : α → Bool) (o : Option α) (h : (o.filter p).isSome) :
+    o.isSome := by
+  cases o <;> simp at h ⊢
+
+@[simp] theorem filter_eq_none (p : α → Bool) :
+    Option.filter p o = none ↔ o = none ∨ ∀ a, a ∈ o → ¬ p a := by
+  cases o <;> simp [filter_some]
+
 @[simp] theorem all_guard (p : α → Prop) [DecidablePred p] (a : α) :
     Option.all q (guard p a) = (!p a || q a) := by
   simp only [guard]
@@ -361,6 +373,20 @@ theorem or_of_isNone {o o' : Option α} (h : o.isNone) : o.or o' = o' := by
   match o, h with
   | none, _ => simp
 
+/-! ### beq -/
+
+section beq
+
+variable [BEq α]
+
+@[simp] theorem none_beq_none : ((none : Option α) == none) = true := rfl
+@[simp] theorem none_beq_some (a : α) : ((none : Option α) == some a) = false := rfl
+@[simp] theorem some_beq_none (a : α) : ((some a : Option α) == none) = false := rfl
+@[simp] theorem some_beq_some {a b : α} : (some a == some b) = (a == b) := rfl
+
+end beq
+
+/-! ### ite -/
 section ite
 
 @[simp] theorem isSome_dite {p : Prop} [Decidable p] {b : p → β} :