Proof of the multinomial theorem.

multinomial.lean

import algebra.big_operators.fin
import algebra.big_operators.order
import data.nat.choose.basic
import data.finset.sym
import data.finsupp.multiset
import data.fin.vec_notation
import data.nat.choose.sum

open_locale nat big_operators

variables {α β : Type*} (s : finset α) (f : α → ℕ)

def multinomial : ℕ := (∑ i in s, f i)! / ∏ i in s, (f i)!

def finsupp.multinomial (f : α →₀ ℕ) : ℕ := (f.sum $ λ _, id)! / f.prod (λ _ n, n!)

lemma finsupp.multinomial_eq (f : α →₀ ℕ) : f.multinomial = multinomial f.support f := rfl

noncomputable def multiset.multinomial (m : multiset α) : ℕ := m.to_finsupp.multinomial

lemma multinomial_insert [decidable_eq α] {a : α} (h : a ∉ s) :
  multinomial (insert a s) f = (f a + s.sum f).choose (f a) * multinomial s f :=
sorry -- from PR #16170

lemma multinomial_congr {f g : α → ℕ} (h : ∀ a ∈ s, f a = g a) :
  multinomial s f = multinomial s g :=
begin
  simp only [multinomial], congr' 1, 
  { rw finset.sum_congr rfl h },
  { exact finset.prod_congr rfl (λ a ha, by rw h a ha) },
end

def sym.fill (a : α) {n : ℕ} (m : Σ i : fin (n + 1), sym α (n - i)) : sym α n :=
sym.mk (m.1.1 • {a} + m.2) begin
  erw [multiset.card_add, m.2.2, multiset.card_nsmul, multiset.card_singleton, mul_one],
  exact nat.add_sub_of_le (nat.lt_succ_iff.1 m.1.2),
end

def sym.filter_ne [decidable_eq α] (a : α) {n : ℕ}
  (m : sym α n) : Σ i : fin (n + 1), sym α (n - i) :=
⟨⟨m.1.count a, (multiset.count_le_card _ _).trans_lt $ by rw [m.2, nat.lt_succ_iff]⟩,
  sym.mk (m.1.filter ((≠) a)) $ eq_tsub_of_add_eq begin
    conv_rhs { rw [← m.2, multiset.card_eq_countp_add_countp ((=) a), add_comm] },
    rw multiset.countp_eq_card_filter, refl,
  end⟩

lemma sym.sigma_ext {n : ℕ} (m₁ m₂ : Σ i : fin (n + 1), sym α (n - i))
  (h : m₁.2.1 = m₂.2.1) : m₁ = m₂ :=
sigma.subtype_ext (fin.ext $ begin
  have h₁ := nat.sub_sub_self (nat.lt_succ_iff.1 m₁.1.2),
  have h₂ := nat.sub_sub_self (nat.lt_succ_iff.1 m₂.1.2),
  dsimp only [fin.val_eq_coe] at h₁ h₂,
  rw [← h₁, ← h₂, ← m₁.2.2, ← m₂.2.2, h],
end) h

lemma finsupp.multinomial_update (a : α) (f : α →₀ ℕ) :
  f.multinomial = (f.sum $ λ _, id).choose (f a) * (f.update a 0).multinomial :=
begin
  simp only [finsupp.multinomial_eq],
  classical,
  by_cases a ∈ f.support,
  { rw [← finset.insert_erase h, multinomial_insert _ f (finset.not_mem_erase a _),
      finset.add_sum_erase _ f h, finsupp.support_update_zero], congr' 1,
    exact multinomial_congr _ (λ _ h, (function.update_noteq (finset.mem_erase.1 h).1 0 f).symm) },
  rw finsupp.not_mem_support_iff at h,
  rw [h, nat.choose_zero_right, one_mul, ← h, finsupp.update_self],
end

lemma multiset.multinomial_filter_ne [decidable_eq α] (a : α) (m : multiset α) :
  m.multinomial = m.card.choose (m.count a) * (m.filter ((≠) a)).multinomial :=
begin
  dsimp only [multiset.multinomial],
  convert finsupp.multinomial_update a _,
  { rw [← finsupp.card_to_multiset, m.to_finsupp_to_multiset] },
  { ext1 a', rw [multiset.to_finsupp_apply, multiset.count_filter, finsupp.coe_update],
    split_ifs,
    { rw [function.update_noteq h.symm, multiset.to_finsupp_apply] },
    { rw [not_ne_iff.1 h, function.update_same] } },
end

def multinomial_theorem [decidable_eq α] {R : Type*} [comm_semiring R] (x : α → R) :
  ∀ n, (s.sum x) ^ n = ∑ k in s.sym n, k.val.multinomial * (k.val.map x).prod :=
begin
  induction s using finset.induction with a s ha ih,
  { rw finset.sum_empty,
    rintro (_ | n),
    { rw [pow_zero, finset.sum_unique_nonempty],
      { convert (one_mul _).symm, apply nat.cast_one },
      { apply finset.univ_nonempty } },
    { rw [pow_succ, zero_mul, finset.sym_empty, finset.sum_empty] } },
  intro n,
  rw [finset.sum_insert ha, add_pow, finset.sum_range],
  simp_rw [ih, finset.mul_sum, finset.sum_mul, finset.sum_sigma'],
  refine (finset.sum_bij (λ m _, sym.filter_ne a m)
    (λ m hm, _) (λ m hm, _) (λ m₁ m₂ h₁ h₂ he, _) (λ m hm, _)).symm,
  { rw finset.mem_sigma,
    rw finset.mem_sym_iff at hm ⊢,
    dsimp only [sym.filter_ne, sym.mem_mk],
    refine ⟨finset.mem_univ _, λ a', _⟩,
    rw multiset.mem_filter,
    exact λ h, finset.mem_of_mem_insert_of_ne (hm a' h.1) h.2.symm },
  { rw [m.1.multinomial_filter_ne a],
    dsimp only [sym.filter_ne, fin.coe_mk],
    conv in (m.1.map _) { rw [← m.1.filter_add_not ((=) a), multiset.map_add] },
    rw [multiset.prod_add, m.1.filter_eq, multiset.map_repeat, multiset.prod_repeat, m.2],
    rw [nat.cast_mul, mul_assoc, mul_comm],
    congr' 1, apply mul_left_comm },
  { replace he := sigma.subtype_ext_iff.1 he,
    dsimp only [sym.filter_ne, subtype.coe_mk] at he,
    simp only [fin.mk.inj_iff] at he,
    ext a', obtain rfl | h := eq_or_ne a a', { exact he.1 },
    erw [← multiset.count_filter_of_pos h, he.2, multiset.count_filter_of_pos h], refl },
  { rw [finset.mem_sigma, finset.mem_sym_iff] at hm,
    refine ⟨sym.fill a m, finset.mem_sym_iff.2 (λ a' h', finset.mem_insert.2 _), _⟩, 
    { rw [sym.fill, sym.mem_mk, multiset.mem_add] at h',
      exact h'.imp (λ h, multiset.mem_singleton.1 (multiset.mem_of_mem_nsmul h)) (λ h, hm.2 a' h) },
    apply sym.sigma_ext, ext1 a',
    dsimp only [sym.filter_ne, sym.fill],
    rw [multiset.count_filter], split_ifs,
    { rw [multiset.count_add, multiset.count_nsmul, multiset.count_singleton, if_neg h.symm],
      rw [mul_zero, zero_add], refl },
    { exact multiset.count_eq_zero.2 (λ h', ha $ (not_ne_iff.1 h).symm ▸ hm.2 a' h') } },
end

-- not used
lemma multinomial_empty : multinomial ∅ f = 1 := rfl
def nil_coe : (coe (@sym.nil α)) = (0 : multiset α) := rfl

-- not used
lemma multiset.to_finsupp_map (f : α → β) (m : multiset α) :
  (m.map f).to_finsupp = m.to_finsupp.map_domain f :=
by conv_lhs
  { rw [← m.to_finsupp_to_multiset, finsupp.to_multiset_map, finsupp.to_multiset_to_finsupp] }

-- not used
lemma multiset.map_multinomial (f : α ↪ β) (m : multiset α) : 
  (m.map f).multinomial = m.multinomial :=
begin
  simp only [multiset.multinomial, finsupp.multinomial, m.to_finsupp_map f],
  congr,
  exacts [finsupp.sum_map_domain_index_inj f.inj', finsupp.prod_map_domain_index_inj f.inj'],
end

-- not used
lemma multinomial_insert_zero [decidable_eq α] (a : α) (h₀ : f a = 0) :
  multinomial (insert a s) f = multinomial s f :=
begin
  by_cases a ∈ s, { rw finset.insert_eq_of_mem h },
  rw [multinomial_insert s f h, h₀, nat.choose_zero_right, one_mul],
end