From 0a1d2908aa37f534bfda61834fe7a7dee9449512 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 22 Nov 2017 08:02:13 -0800
Subject: [PATCH] feat(library/data/rbtree/main): rbtrees are balanced

---
 library/data/rbtree/main.lean       | 17 +++++++++++++++++
 library/init/data/rbtree/basic.lean |  3 +++
 2 files changed, 20 insertions(+)

diff --git a/library/data/rbtree/main.lean b/library/data/rbtree/main.lean
index c564ad75ff..62d6d3a957 100644
--- a/library/data/rbtree/main.lean
+++ b/library/data/rbtree/main.lean
@@ -19,11 +19,28 @@ begin
   { subst n', apply is_searchable_insert, assumption }
 end
 
+open color
+
+lemma is_red_black_of_well_formed {t : rbnode α} : t.well_formed lt → ∃ c n, is_red_black t c n :=
+begin
+  intro h, induction h,
+  { existsi black, existsi 0, constructor },
+  { cases ih_1 with c ih, cases ih with n ih, subst n', apply insert_is_red_black, assumption }
+end
+
 end rbnode
 
 namespace rbtree
 variables {α : Type u} {lt : α → α → Prop}
 
+lemma balanced (t : rbtree α lt) : 2 * t.depth min + 1 ≥ t.depth max :=
+begin
+  cases t with n p, simp only [depth],
+  have := rbnode.is_red_black_of_well_formed p,
+  cases this with _ this, cases this with _ this,
+  apply rbnode.balanced, assumption
+end
+
 lemma not_mem_mk_rbtree : ∀ (a : α), a ∉ mk_rbtree α lt :=
 by simp [has_mem.mem, rbtree.mem, rbnode.mem, mk_rbtree]
 
diff --git a/library/init/data/rbtree/basic.lean b/library/init/data/rbtree/basic.lean
index c907745fa3..5565756368 100644
--- a/library/init/data/rbtree/basic.lean
+++ b/library/init/data/rbtree/basic.lean
@@ -173,6 +173,9 @@ instance : has_mem α (rbtree α lt) :=
 def mem_exact (a : α) (t : rbtree α lt) : Prop :=
 rbnode.mem_exact a t.val
 
+def depth (f : nat → nat → nat) (t : rbtree α lt) : nat :=
+t.val.depth f
+
 def fold (f : α → β → β) : rbtree α lt → β →  β
 | ⟨t, _⟩ b := t.fold f b