feat: List/Array/Vector.count_replace lemmas

src/Init/Data/Array/Count.lean

@@ -288,6 +288,17 @@ theorem count_flatMap {α} [BEq β] {xs : Array α} {f : α → Array β} {x :
   rcases xs with ⟨xs⟩
   simp [List.count_flatMap, countP_flatMap, Function.comp_def]
 
+theorem countP_replace {a b : α} {xs : Array α} {p : α → Bool} :
+    (xs.replace a b).countP p =
+      if xs.contains a then xs.countP p + (if p b then 1 else 0) - (if p a then 1 else 0) else xs.countP p := by
+  rcases xs with ⟨xs⟩
+  simp [List.countP_replace]
+
+theorem count_replace {a b c : α} {xs : Array α} :
+    (xs.replace a b).count c =
+      if xs.contains a then xs.count c + (if b == c then 1 else 0) - (if a == c then 1 else 0) else xs.count c := by
+  simp [count_eq_countP, countP_replace]
+
 -- FIXME these theorems can be restored once `List.erase` and `Array.erase` have been related.
 
 -- theorem count_erase (a b : α) (l : Array α) : count a (l.erase b) = count a l - if b == a then 1 else 0 := by

src/Init/Data/List/Nat/Count.lean

@@ -31,6 +31,33 @@ 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]
 
+theorem countP_replace [BEq α] [LawfulBEq α] {a b : α} {l : List α} {p : α → Bool} :
+    (l.replace a b).countP p =
+      if l.contains a then l.countP p + (if p b then 1 else 0) - (if p a then 1 else 0) else l.countP p := by
+  induction l with
+  | nil => simp
+  | cons x l ih =>
+    simp [replace_cons]
+    split <;> rename_i h
+    · simp at h
+      simp [h, ih, countP_cons]
+      omega
+    · simp only [beq_eq_false_iff_ne, ne_eq] at h
+      simp only [countP_cons, ih, contains_eq_mem, decide_eq_true_eq, mem_cons, h, false_or]
+      split <;> rename_i h'
+      · by_cases h'' : p a
+        · have : countP p l > 0 := countP_pos_iff.mpr ⟨a, h', h''⟩
+          simp [h'']
+          omega
+        · simp [h'']
+          omega
+      · omega
+
+theorem count_replace [BEq α] [LawfulBEq α] {a b c : α} {l : List α} :
+    (l.replace a b).count c =
+      if l.contains a then l.count c + (if b == c then 1 else 0) - (if a == c then 1 else 0) else l.count c := by
+  simp [count_eq_countP, countP_replace]
+
 /--
 The number of elements satisfying a predicate in a sublist is at least the number of elements satisfying the predicate in the list,
 minus the difference in the lengths.

src/Init/Data/Vector/Count.lean

@@ -239,4 +239,15 @@ theorem count_flatMap {α} [BEq β] {xs : Vector α n} {f : α → Vector β m}
   rcases xs with ⟨xs, rfl⟩
   simp [Array.count_flatMap, Function.comp_def]
 
+theorem countP_replace {a b : α} {xs : Vector α n} {p : α → Bool} :
+    (xs.replace a b).countP p =
+      if xs.contains a then xs.countP p + (if p b then 1 else 0) - (if p a then 1 else 0) else xs.countP p := by
+  rcases xs with ⟨xs, rfl⟩
+  simp [Array.countP_replace]
+
+theorem count_replace {a b c : α} {xs : Vector α n} :
+    (xs.replace a b).count c =
+      if xs.contains a then xs.count c + (if b == c then 1 else 0) - (if a == c then 1 else 0) else xs.count c := by
+  simp [count_eq_countP, countP_replace]
+
 end count