-
Notifications
You must be signed in to change notification settings - Fork 307
/
Tree.lean
167 lines (135 loc) · 5.68 KB
/
Tree.lean
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
/-
Copyright (c) 2019 mathlib community. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Mario Carneiro, Wojciech Nawrocki
-/
import Mathlib.Data.Num.Basic
import Mathlib.Init.Data.Ordering.Basic
import Mathlib.Init.Order.LinearOrder
import Mathlib.Util.CompileInductive
import Std.Data.RBMap.Basic
#align_import data.tree from "leanprover-community/mathlib"@"ed989ff568099019c6533a4d94b27d852a5710d8"
/-!
# Binary tree
Provides binary tree storage for values of any type, with O(lg n) retrieval.
See also `Lean.Data.RBTree` for red-black trees - this version allows more operations
to be defined and is better suited for in-kernel computation.
We also specialize for `Tree Unit`, which is a binary tree without any
additional data. We provide the notation `a △ b` for making a `Tree Unit` with children
`a` and `b`.
## References
<https://leanprover-community.github.io/archive/stream/113488-general/topic/tactic.20question.html>
-/
/-- A binary tree with values stored in non-leaf nodes. -/
inductive Tree.{u} (α : Type u) : Type u
| nil : Tree α
| node : α → Tree α → Tree α → Tree α
deriving DecidableEq, Repr -- Porting note: Removed `has_reflect`, added `Repr`.
#align tree Tree
namespace Tree
universe u
variable {α : Type u}
-- Porting note: replaced with `deriving Repr` which builds a better instance anyway
#noalign tree.repr
instance : Inhabited (Tree α) :=
⟨nil⟩
open Std (RBNode)
/-- Makes a `Tree α` out of a red-black tree. -/
def ofRBNode : RBNode α → Tree α
| RBNode.nil => nil
| RBNode.node _color l key r => node key (ofRBNode l) (ofRBNode r)
#align tree.of_rbnode Tree.ofRBNode
/-- Finds the index of an element in the tree assuming the tree has been
constructed according to the provided decidable order on its elements.
If it hasn't, the result will be incorrect. If it has, but the element
is not in the tree, returns none. -/
def indexOf (lt : α → α → Prop) [DecidableRel lt] (x : α) : Tree α → Option PosNum
| nil => none
| node a t₁ t₂ =>
match cmpUsing lt x a with
| Ordering.lt => PosNum.bit0 <$> indexOf lt x t₁
| Ordering.eq => some PosNum.one
| Ordering.gt => PosNum.bit1 <$> indexOf lt x t₂
#align tree.index_of Tree.indexOf
/-- Retrieves an element uniquely determined by a `PosNum` from the tree,
taking the following path to get to the element:
- `bit0` - go to left child
- `bit1` - go to right child
- `PosNum.one` - retrieve from here -/
def get : PosNum → Tree α → Option α
| _, nil => none
| PosNum.one, node a _t₁ _t₂ => some a
| PosNum.bit0 n, node _a t₁ _t₂ => t₁.get n
| PosNum.bit1 n, node _a _t₁ t₂ => t₂.get n
#align tree.get Tree.get
/-- Retrieves an element from the tree, or the provided default value
if the index is invalid. See `Tree.get`. -/
def getOrElse (n : PosNum) (t : Tree α) (v : α) : α :=
(t.get n).getD v
#align tree.get_or_else Tree.getOrElse
/-- Apply a function to each value in the tree. This is the `map` function for the `Tree` functor.
TODO: implement `Traversable Tree`. -/
def map {β} (f : α → β) : Tree α → Tree β
| nil => nil
| node a l r => node (f a) (map f l) (map f r)
#align tree.map Tree.map
/-- The number of internal nodes (i.e. not including leaves) of a binary tree -/
@[simp]
def numNodes : Tree α → ℕ
| nil => 0
| node _ a b => a.numNodes + b.numNodes + 1
#align tree.num_nodes Tree.numNodes
/-- The number of leaves of a binary tree -/
@[simp]
def numLeaves : Tree α → ℕ
| nil => 1
| node _ a b => a.numLeaves + b.numLeaves
#align tree.num_leaves Tree.numLeaves
/-- The height - length of the longest path from the root - of a binary tree -/
@[simp]
def height : Tree α → ℕ
| nil => 0
| node _ a b => max a.height b.height + 1
#align tree.height Tree.height
theorem numLeaves_eq_numNodes_succ (x : Tree α) : x.numLeaves = x.numNodes + 1 := by
induction x <;> simp [*, Nat.add_comm, Nat.add_assoc, Nat.add_left_comm]
#align tree.num_leaves_eq_num_nodes_succ Tree.numLeaves_eq_numNodes_succ
theorem numLeaves_pos (x : Tree α) : 0 < x.numLeaves := by
rw [numLeaves_eq_numNodes_succ]
exact x.numNodes.zero_lt_succ
#align tree.num_leaves_pos Tree.numLeaves_pos
theorem height_le_numNodes : ∀ x : Tree α, x.height ≤ x.numNodes
| nil => le_rfl
| node _ a b =>
Nat.succ_le_succ
(max_le (_root_.trans a.height_le_numNodes <| a.numNodes.le_add_right _)
(_root_.trans b.height_le_numNodes <| b.numNodes.le_add_left _))
#align tree.height_le_num_nodes Tree.height_le_numNodes
/-- The left child of the tree, or `nil` if the tree is `nil` -/
@[simp]
def left : Tree α → Tree α
| nil => nil
| node _ l _r => l
#align tree.left Tree.left
/-- The right child of the tree, or `nil` if the tree is `nil` -/
@[simp]
def right : Tree α → Tree α
| nil => nil
| node _ _l r => r
#align tree.right Tree.right
-- Notation for making a node with `Unit` data
scoped infixr:65 " △ " => Tree.node ()
-- Porting note: workaround for leanprover/lean4#2049
compile_inductive% Tree
@[elab_as_elim]
def unitRecOn {motive : Tree Unit → Sort*} (t : Tree Unit) (base : motive nil)
(ind : ∀ x y, motive x → motive y → motive (x △ y)) : motive t :=
-- Porting note: Old proof was `t.recOn base fun u => u.recOn ind` but
-- structure eta makes it unnecessary (https://github.com/leanprover/lean4/issues/777).
t.recOn base fun _u => ind
#align tree.unit_rec_on Tree.unitRecOn
theorem left_node_right_eq_self : ∀ {x : Tree Unit} (_hx : x ≠ nil), x.left △ x.right = x
| nil, h => by trivial
| node a l r, _ => rfl -- Porting note: `a △ b` no longer works in pattern matching
#align tree.left_node_right_eq_self Tree.left_node_right_eq_self
end Tree