feat(data/list/perm): nodup_permutations (#8616)

A simpler version of #8494

TODO: `nodup s.permutations ↔ nodup s`
TODO: `count s s.permutations = (zip_with count s s.tails).prod`



Co-authored-by: Yakov Pechersky <pechersky@users.noreply.github.com>

src/data/list/basic.lean

@@ -1409,10 +1409,13 @@ lemma mem_or_eq_of_mem_update_nth : ∀ {l : list α} {n : ℕ} {a b : α}
 section insert_nth
 variable {a : α}
 
-@[simp] lemma insert_nth_nil (a : α) : insert_nth 0 a [] = [a] := rfl
+@[simp] lemma insert_nth_zero (s : list α) (x : α) : insert_nth 0 x s = x :: s := rfl
 
 @[simp] lemma insert_nth_succ_nil (n : ℕ) (a : α) : insert_nth (n + 1) a [] = [] := rfl
 
+@[simp] lemma insert_nth_succ_cons (s : list α) (hd x : α) (n : ℕ) :
+  insert_nth (n + 1) x (hd :: s) = hd :: (insert_nth n x s) := rfl
+
 lemma length_insert_nth : ∀n as, n ≤ length as → length (insert_nth n a as) = length as + 1
 | 0     as       h := rfl
 | (n+1) []       h := (nat.not_succ_le_zero _ h).elim
@@ -1459,6 +1462,119 @@ lemma mem_insert_nth {a b : α} : ∀ {n : ℕ} {l : list α} (hi : n ≤ l.leng
     ← or.assoc, or_comm (a = a'), or.assoc]
 end
 
+lemma inj_on_insert_nth_index_of_not_mem (l : list α) (x : α) (hx : x ∉ l) :
+  set.inj_on (λ k, insert_nth k x l) {n | n ≤ l.length} :=
+begin
+  induction l with hd tl IH,
+  { intros n hn m hm h,
+    simp only [set.mem_singleton_iff, set.set_of_eq_eq_singleton, length, nonpos_iff_eq_zero]
+      at hn hm,
+    simp [hn, hm] },
+  { intros n hn m hm h,
+    simp only [length, set.mem_set_of_eq] at hn hm,
+    simp only [mem_cons_iff, not_or_distrib] at hx,
+    cases n;
+    cases m,
+    { refl },
+    { simpa [hx.left] using h },
+    { simpa [ne.symm hx.left] using h },
+    { simp only [true_and, eq_self_iff_true, insert_nth_succ_cons] at h,
+      rw nat.succ_inj',
+      refine IH hx.right _ _ h,
+      { simpa [nat.succ_le_succ_iff] using hn },
+      { simpa [nat.succ_le_succ_iff] using hm } } }
+end
+
+lemma insert_nth_of_length_lt (l : list α) (x : α) (n : ℕ) (h : l.length < n) :
+  insert_nth n x l = l :=
+begin
+  induction l with hd tl IH generalizing n,
+  { cases n,
+    { simpa using h },
+    { simp } },
+  { cases n,
+    { simpa using h },
+    { simp only [nat.succ_lt_succ_iff, length] at h,
+      simpa using IH _ h } }
+end
+
+@[simp] lemma insert_nth_length_self (l : list α) (x : α) :
+  insert_nth l.length x l = l ++ [x] :=
+begin
+  induction l with hd tl IH,
+  { simp },
+  { simpa using IH }
+end
+
+lemma length_le_length_insert_nth (l : list α) (x : α) (n : ℕ) :
+  l.length ≤ (insert_nth n x l).length :=
+begin
+  cases le_or_lt n l.length with hn hn,
+  { rw length_insert_nth _ _ hn,
+    exact (nat.lt_succ_self _).le },
+  { rw insert_nth_of_length_lt _ _ _ hn }
+end
+
+lemma length_insert_nth_le_succ (l : list α) (x : α) (n : ℕ) :
+  (insert_nth n x l).length ≤ l.length + 1 :=
+begin
+  cases le_or_lt n l.length with hn hn,
+  { rw length_insert_nth _ _ hn },
+  { rw insert_nth_of_length_lt _ _ _ hn,
+    exact (nat.lt_succ_self _).le }
+end
+
+lemma nth_le_insert_nth_of_lt (l : list α) (x : α) (n k : ℕ) (hn : k < n)
+  (hk : k < l.length)
+  (hk' : k < (insert_nth n x l).length := hk.trans_le (length_le_length_insert_nth _ _ _)):
+  (insert_nth n x l).nth_le k hk' = l.nth_le k hk :=
+begin
+  induction n with n IH generalizing k l,
+  { simpa using hn },
+  { cases l with hd tl,
+    { simp },
+    { cases k,
+      { simp },
+      { rw nat.succ_lt_succ_iff at hn,
+        simpa using IH _ _ hn _ } } }
+end
+
+@[simp] lemma nth_le_insert_nth_self (l : list α) (x : α) (n : ℕ)
+  (hn : n ≤ l.length) (hn' : n < (insert_nth n x l).length :=
+    by rwa [length_insert_nth _ _ hn, nat.lt_succ_iff]) :
+  (insert_nth n x l).nth_le n hn' = x :=
+begin
+  induction l with hd tl IH generalizing n,
+  { simp only [length, nonpos_iff_eq_zero] at hn,
+    simp [hn] },
+  { cases n,
+    { simp },
+    { simp only [nat.succ_le_succ_iff, length] at hn,
+      simpa using IH _ hn } }
+end
+
+lemma nth_le_insert_nth_add_succ (l : list α) (x : α) (n k : ℕ)
+  (hk' : n + k < l.length)
+  (hk : n + k + 1 < (insert_nth n x l).length :=
+    by rwa [length_insert_nth _ _ (le_self_add.trans hk'.le), nat.succ_lt_succ_iff]) :
+  (insert_nth n x l).nth_le (n + k + 1) hk = nth_le l (n + k) hk' :=
+begin
+  induction l with hd tl IH generalizing n k,
+  { simpa using hk' },
+  { cases n,
+    { simpa },
+    { simpa [succ_add] using IH _ _ _ } }
+end
+
+lemma insert_nth_injective (n : ℕ) (x : α) : function.injective (insert_nth n x) :=
+begin
+  induction n with n IH,
+  { have : insert_nth 0 x = cons x := funext (λ _, rfl),
+    simp [this] },
+  { rintros (_|⟨a, as⟩) (_|⟨b, bs⟩) h;
+    simpa [IH.eq_iff] using h <|> refl }
+end
+
 end insert_nth
 
 /-! ### map -/
@@ -3452,6 +3568,26 @@ eq_of_sublist_of_length_eq (le_count_iff_repeat_sublist.mp (le_refl (count a l))
 by simp only [count, countp_filter]; congr; exact
 set.ext (λ b, and_iff_left_of_imp (λ e, e ▸ h))
 
+lemma count_bind {α β} [decidable_eq β] (l : list α) (f : α → list β) (x : β)  :
+  count x (l.bind f) = sum (map (count x ∘ f) l) :=
+begin
+  induction l with hd tl IH,
+  { simp },
+  { simpa }
+end
+
+@[simp] lemma count_map_map {α β} [decidable_eq α] [decidable_eq β] (l : list α) (f : α → β)
+  (hf : function.injective f) (x : α) :
+  count (f x) (map f l) = count x l :=
+begin
+  induction l with y l IH generalizing x,
+  { simp },
+  { rw map_cons,
+    by_cases h : x = y,
+    { simpa [h] using IH _ },
+    { simpa [h, hf.ne h] using IH _ } }
+end
+
 end count
 
 /-! ### prefix, suffix, infix -/

src/data/list/perm.lean

@@ -1168,6 +1168,152 @@ theorem perm.permutations {s t : list α} (h : s ~ t) : permutations s ~ permuta
 @[simp] theorem perm_permutations'_iff {s t : list α} : permutations' s ~ permutations' t ↔ s ~ t :=
 ⟨λ h, mem_permutations'.1 $ h.mem_iff.1 $ mem_permutations'.2 (perm.refl _), perm.permutations'⟩
 
+lemma nth_le_permutations'_aux (s : list α) (x : α) (n : ℕ)
+  (hn : n < length (permutations'_aux x s)) :
+  (permutations'_aux x s).nth_le n hn = s.insert_nth n x :=
+begin
+  induction s with y s IH generalizing n,
+  { simp only [length, permutations'_aux, nat.lt_one_iff] at hn,
+    simp [hn] },
+  { cases n,
+    { simp },
+    { simpa using IH _ _ } }
+end
+
+lemma count_permutations'_aux_self [decidable_eq α] (l : list α) (x : α) :
+  count (x :: l) (permutations'_aux x l) = length (take_while ((=) x) l) + 1 :=
+begin
+  induction l with y l IH generalizing x,
+  { simp [take_while], },
+  { rw [permutations'_aux, count_cons_self],
+    by_cases hx : x = y,
+    { subst hx,
+      simpa [take_while, nat.succ_inj'] using IH _ },
+    { rw take_while,
+      rw if_neg hx,
+      cases permutations'_aux x l with a as,
+      { simp },
+      { rw [count_eq_zero_of_not_mem, length, zero_add],
+        simp [hx, ne.symm hx] } } }
+end
+
+@[simp] lemma length_permutations'_aux (s : list α) (x : α) :
+  length (permutations'_aux x s) = length s + 1 :=
+begin
+  induction s with y s IH,
+  { simp },
+  { simpa using IH }
+end
+
+@[simp] lemma permutations'_aux_nth_le_zero (s : list α) (x : α)
+  (hn : 0 < length (permutations'_aux x s) := by simp) :
+  (permutations'_aux x s).nth_le 0 hn = x :: s :=
+nth_le_permutations'_aux _ _ _ _
+
+lemma injective_permutations'_aux (x : α) : function.injective (permutations'_aux x) :=
+begin
+  intros s t h,
+  apply insert_nth_injective s.length x,
+  have hl : s.length = t.length := by simpa using congr_arg length h,
+  rw [←nth_le_permutations'_aux s x s.length (by simp),
+      ←nth_le_permutations'_aux t x s.length (by simp [hl])],
+  simp [h, hl]
+end
+
+lemma nodup_permutations'_aux_of_not_mem (s : list α) (x : α) (hx : x ∉ s) :
+  nodup (permutations'_aux x s) :=
+begin
+  induction s with y s IH,
+  { simp },
+  { simp only [not_or_distrib, mem_cons_iff] at hx,
+    simp only [not_and, exists_eq_right_right, mem_map, permutations'_aux, nodup_cons],
+    refine ⟨λ _, ne.symm hx.left, _⟩,
+    rw nodup_map_iff,
+    { exact IH hx.right },
+    { simp } }
+end
+
+lemma nodup_permutations'_aux_iff {s : list α} {x : α} :
+  nodup (permutations'_aux x s) ↔ x ∉ s :=
+begin
+  refine ⟨λ h, _, nodup_permutations'_aux_of_not_mem _ _⟩,
+  intro H,
+  obtain ⟨k, hk, hk'⟩ := nth_le_of_mem H,
+  rw nodup_iff_nth_le_inj at h,
+  suffices : k = k + 1,
+  { simpa using this },
+  refine h k (k + 1) _ _ _,
+  { simpa [nat.lt_succ_iff] using hk.le },
+  { simpa using hk },
+  rw [nth_le_permutations'_aux, nth_le_permutations'_aux],
+  have hl : length (insert_nth k x s) = length (insert_nth (k + 1) x s),
+  { rw [length_insert_nth _ _ hk.le, length_insert_nth _ _ (nat.succ_le_of_lt hk)] },
+  refine ext_le hl (λ n hn hn', _),
+  rcases lt_trichotomy n k with H|rfl|H,
+  { rw [nth_le_insert_nth_of_lt _ _ _ _ H (H.trans hk),
+        nth_le_insert_nth_of_lt _ _ _ _ (H.trans (nat.lt_succ_self _))] },
+  { rw [nth_le_insert_nth_self _ _ _ hk.le,
+        nth_le_insert_nth_of_lt _ _ _ _ (nat.lt_succ_self _) hk, hk'] },
+  { rcases (nat.succ_le_of_lt H).eq_or_lt with rfl|H',
+    { rw [nth_le_insert_nth_self _ _ _ (nat.succ_le_of_lt hk)],
+      convert hk' using 1,
+      convert nth_le_insert_nth_add_succ _ _ _ 0 _,
+      simpa using hk },
+    { obtain ⟨m, rfl⟩ := nat.exists_eq_add_of_lt H',
+      rw [length_insert_nth _ _ hk.le, nat.succ_lt_succ_iff, nat.succ_add] at hn,
+      rw nth_le_insert_nth_add_succ,
+      convert nth_le_insert_nth_add_succ s x k m.succ _ using 2,
+      { simp [nat.add_succ, nat.succ_add] },
+      { simp [add_left_comm, add_comm] },
+      { simpa [nat.add_succ] using hn },
+      { simpa [nat.succ_add] using hn } } }
+end
+
+lemma nodup_permutations (s : list α) (hs : nodup s) :
+  nodup s.permutations :=
+begin
+  rw (permutations_perm_permutations' s).nodup_iff,
+  induction hs with x l h h' IH,
+  { simp },
+  { rw [permutations'],
+    rw nodup_bind,
+    split,
+    { intros ys hy,
+      rw mem_permutations' at hy,
+      rw [nodup_permutations'_aux_iff, hy.mem_iff],
+      exact λ H, h x H rfl },
+    { refine IH.pairwise_of_forall_ne (λ as ha bs hb H, _),
+      rw disjoint_iff_ne,
+      rintro a ha' b hb' rfl,
+      obtain ⟨n, hn, hn'⟩ := nth_le_of_mem ha',
+      obtain ⟨m, hm, hm'⟩ := nth_le_of_mem hb',
+      rw mem_permutations' at ha hb,
+      have hl : as.length = bs.length := (ha.trans hb.symm).length_eq,
+      simp only [nat.lt_succ_iff, length_permutations'_aux] at hn hm,
+      rw nth_le_permutations'_aux at hn' hm',
+      have hx : nth_le (insert_nth n x as) m
+        (by rwa [length_insert_nth _ _ hn, nat.lt_succ_iff, hl]) = x,
+      { simp [hn', ←hm', hm] },
+      have hx' : nth_le (insert_nth m x bs) n
+        (by rwa [length_insert_nth _ _ hm, nat.lt_succ_iff, ←hl]) = x,
+      { simp [hm', ←hn', hn] },
+      rcases lt_trichotomy n m with ht|ht|ht,
+      { suffices : x ∈ bs,
+        { exact h x (hb.subset this) rfl },
+        rw [←hx', nth_le_insert_nth_of_lt _ _ _ _ ht (ht.trans_le hm)],
+        exact nth_le_mem _ _ _ },
+      { simp only [ht] at hm' hn',
+        rw ←hm' at hn',
+        exact H (insert_nth_injective _ _ hn') },
+      { suffices : x ∈ as,
+        { exact h x (ha.subset this) rfl },
+        rw [←hx, nth_le_insert_nth_of_lt _ _ _ _ ht (ht.trans_le hn)],
+        exact nth_le_mem _ _ _ } } }
+end
+
+-- TODO: `nodup s.permutations ↔ nodup s`
+-- TODO: `count s s.permutations = (zip_with count s s.tails).prod`
+
 end permutations
 
 end list