feat: splay trees improve

Cslib/Foundations/Data/BinaryTree.lean

@@ -7,71 +7,38 @@ Authors: Anton Kovsharov, Antoine du Fresne von Hohenesche,
 
 module
 
+public import Mathlib.Data.Tree.Basic
+public import Mathlib.Data.Tree.Traversable
+public import Mathlib.Control.Fold.Traversable
 public import Mathlib.Combinatorics.SimpleGraph.Basic
 public import Mathlib.Combinatorics.SimpleGraph.Metric
 
-@[expose] public section
-
 /-!
-# Binary Tree
+# Tree
 
-In this file we introduce the `BinaryTree` data structure and its basic operations.
+In this file we extend the `Tree` data structure from mathlib with basic operations.
 -/
 
-variable {α : Type}
-
-inductive BinaryTree (α : Type) where
-  | empty
-  | node (left : BinaryTree α) (key : α) (right : BinaryTree α)
-  deriving DecidableEq
-
-namespace BinaryTree
-
-/-! ### Core Definitions -/
-section CoreDefs
-
-def left : BinaryTree α → BinaryTree α
-  | .empty => .empty
-  | .node l _ _ => l
-
-def right : BinaryTree α → BinaryTree α
-  | .empty => .empty
-  | .node _ _ r => r
-
-theorem non_empty_exist (s : BinaryTree α) (h : s ≠ .empty) :
-    ∃ A k B, s = .node A k B := by
-  induction s <;> grind
-
-def num_nodes : BinaryTree α → ℕ
-  | .empty => 0
-  | .node l _ r => 1 + num_nodes l + num_nodes r
-
-@[simp] lemma num_nodes_empty : num_nodes (empty : BinaryTree α) = 0 := rfl
+@[expose] public section
 
-@[simp] lemma num_nodes_node (l : BinaryTree α) (k : α) (r : BinaryTree α) :
-    (node l k r).num_nodes = 1 + l.num_nodes + r.num_nodes := rfl
+namespace Cslib
 
-/-- In-order traversal as a list of keys. -/
-def toKeyList : BinaryTree α → List α
-  | .empty => []
-  | .node l k r => l.toKeyList ++ [k] ++ r.toKeyList
+namespace Tree
 
-@[simp] lemma toKeyList_empty : toKeyList (empty : BinaryTree α) = [] := rfl
+open Cslib
 
-@[simp] lemma toKeyList_node (l : BinaryTree α) (k : α) (r : BinaryTree α) :
-    (node l k r).toKeyList = l.toKeyList ++ [k] ++ r.toKeyList := rfl
 
 /-- Number of nodes on the search path for `q` in `t`. Zero on the empty
 tree; on a node this counts the root plus (if `q ≠ k`) the search path
 length in the appropriate subtree. -/
-def search_path_len [LinearOrder α] (t : BinaryTree α) (q : α) : ℕ :=
+def search_path_len [LinearOrder α] (t : Tree α) (q : α) : ℕ :=
   match t with
-  | .empty => 0
-  | .node l key r =>
+  | .nil => 0
+  | .node key l r =>
     if q < key then
-      1 + l.search_path_len q
+      1 + search_path_len l q
     else if key < q then
-      1 + r.search_path_len q
+      1 + search_path_len r q
     else
       1
 
@@ -84,68 +51,68 @@ If `contains t q` is true, then `q` is in `t`; but
 the converse need not necessarily hold true. The
 converse is true for a binary search tree.
 -/
-def contains [LinearOrder α] (t : BinaryTree α) (q : α) : Prop :=
+def containsBST [LinearOrder α] (t : Tree α) (q : α) : Prop :=
   match t with
-  | .empty => False
-  | .node l key r =>
+  | .nil => False
+  | .node key l r =>
     if q < key then
-      l.contains q
+      containsBST l q
     else if key < q then
-      r.contains q
+      containsBST r q
     else
       True
 
-end CoreDefs
-
 
 /-! ### Rotations and Mirroring -/
 section Transformations
 
-def rotateRight : BinaryTree α → BinaryTree α
-  | .node (.node a x b) y c => .node a x (.node b y c)
+def rotateRight : Tree α → Tree α
+  | .node y (.node x a b) c => .node x a (.node y b c)
   | t => t
 
-def rotateLeft : BinaryTree α → BinaryTree α
-  | .node a x (.node b y c) => .node (.node a x b) y c
+def rotateLeft : Tree α → Tree α
+  | .node x a (.node y b c) => .node y (.node x a b) c
   | t => t
 
 /-- Mirror a binary tree: swap every left and right subtree. -/
-def mirror : BinaryTree α → BinaryTree α
-  | .empty => .empty
-  | .node l k r => .node r.mirror k l.mirror
+def mirror : Tree α → Tree α
+  | .nil => .nil
+  | .node k l r => .node k (mirror r) (mirror l)
 
-@[simp] lemma mirror_empty : (empty : BinaryTree α).mirror = .empty := rfl
+@[simp] lemma mirror_empty : (.nil : Tree α).mirror = .nil := rfl
 
-@[simp] lemma mirror_node (l : BinaryTree α) (k : α) (r : BinaryTree α) :
-    (node l k r).mirror = .node r.mirror k l.mirror := rfl
+@[simp] lemma mirror_node (k : α) (l : Tree α) (r : Tree α) :
+    (Tree.node k l r).mirror = .node k r.mirror l.mirror := rfl
 
-@[simp] lemma mirror_mirror (t : BinaryTree α) : t.mirror.mirror = t := by
+
+@[simp] lemma mirror_mirror (t : Tree α) : t.mirror.mirror = t := by
   induction t <;> simp_all
 
-@[simp] lemma num_nodes_mirror (t : BinaryTree α) : t.mirror.num_nodes = t.num_nodes := by
-  induction t <;> simp_all [num_nodes]; omega
+@[simp] lemma size_mirror (t : Tree α) : t.mirror.numNodes = t.numNodes := by
+  induction t <;> simp_all [Tree.numNodes]; omega
 
-@[simp] lemma mirror_rotateRight (t : BinaryTree α) :
+@[simp] lemma mirror_rotateRight (t : Tree α) :
     (rotateRight t).mirror = rotateLeft t.mirror := by
-  cases t; · rfl
-  rename_i l _ _
-  cases l <;> rfl
+  cases t
+  · rfl
+  · rename_i _ l _
+    cases l <;> rfl
 
-@[simp] lemma mirror_rotateLeft (t : BinaryTree α) :
+@[simp] lemma mirror_rotateLeft (t : Tree α) :
     (rotateLeft t).mirror = rotateRight t.mirror := by
   cases t; · rfl
   rename_i _ _ r
   cases r <;> rfl
 
-@[simp] theorem num_nodes_rotateRight (t : BinaryTree α) :
-    (rotateRight t).num_nodes = t.num_nodes := by
-  rcases t with _ | ⟨(_ | ⟨ll, lk, lr⟩), k, r⟩ <;>
+@[simp] theorem size_rotateRight (t : Tree α) :
+    (rotateRight t).numNodes = t.numNodes := by
+  rcases t with _ | ⟨k, (_ | ⟨lk, ll, lr⟩), r⟩ <;>
     simp [rotateRight]; omega
 
-@[simp] theorem num_nodes_rotateLeft (t : BinaryTree α) :
-    (rotateLeft t).num_nodes = t.num_nodes := by
-  have h := num_nodes_rotateRight t.mirror
-  simp only [← mirror_rotateLeft, num_nodes_mirror] at h; exact h
+@[simp] theorem size_rotateLeft (t : Tree α) :
+    (rotateLeft t).numNodes = t.numNodes := by
+  have h := size_rotateRight t.mirror
+  simp only [← mirror_rotateLeft, size_mirror] at h; exact h
 
 end Transformations
 
@@ -154,62 +121,58 @@ end Transformations
 section ContainsLemmas
 
 @[simp] lemma not_contains_empty [LinearOrder α] (q : α) :
-    ¬ (empty : BinaryTree α).contains q := nofun
+    ¬ (Tree.nil : Tree α).containsBST q := nofun
 
-@[simp] lemma contains_node_lt [LinearOrder α] {l : BinaryTree α} {k q : α}
-    {r : BinaryTree α} (h : q < k) :
-    (node l k r).contains q ↔ l.contains q := by
-  simp [contains, h]
+@[simp] lemma contains_node_lt [LinearOrder α] {l : Tree α} {k q : α}
+    {r : Tree α} (h : q < k) :
+    (Tree.node k l r).containsBST q ↔ l.containsBST q := by
+  simp [containsBST, h]
 
-@[simp] lemma contains_node_gt [LinearOrder α] {l : BinaryTree α} {k q : α}
-    {r : BinaryTree α} (h : k < q) :
-    (node l k r).contains q ↔ r.contains q := by
-  simp [contains, h, not_lt_of_gt h]
+@[simp] lemma contains_node_gt [LinearOrder α] {l : Tree α} {k q : α}
+    {r : Tree α} (h : k < q) :
+    (Tree.node k l r).containsBST q ↔ r.containsBST q := by
+  simp [containsBST, h, not_lt_of_gt h]
 
 @[simp] lemma contains_node_not_eq_not_lt [LinearOrder α]
-    {l : BinaryTree α} {k q : α} {r : BinaryTree α}
+    {l : Tree α} {k q : α} {r : Tree α}
     (h1 : ¬ q = k) (h2 : ¬ q < k) :
-    (node l k r).contains q ↔ r.contains q := by
+    (Tree.node k l r).containsBST q ↔ r.containsBST q := by
   have hgt : k < q := lt_of_le_of_ne (Std.not_lt.mp h2) (Ne.symm (Ne.intro h1))
-  simp [contains, hgt, not_lt_of_gt hgt]
+  simp [containsBST, hgt, not_lt_of_gt hgt]
 
 end ContainsLemmas
 
 
 /-! ### Tree Invariants and BST Properties -/
 section Invariants
 
-inductive ForallTree (p : α → Prop) : BinaryTree α → Prop
-  | left : ForallTree p .empty
-  | node l key r :
-     ForallTree p l →
-     p key →
-     ForallTree p r →
-     ForallTree p (.node l key r)
-
-inductive IsBST [LinearOrder α] : BinaryTree α → Prop
-  | left : IsBST .empty
-  | node key l r :
-     ForallTree (fun k => k < key) l →
-     ForallTree (fun k => key < k) r →
-     IsBST l → IsBST r →
-     IsBST (.node l key r)
+
+def contains (t : Tree β) (p : β) : Prop :=
+  match t with
+  | .nil => False
+  | .node k l r => p = k ∨ Tree.contains l p ∨ Tree.contains r p
+
+instance : Membership β (Tree β) where
+  mem := Tree.contains
+
+def IsBST [LinearOrder α] (t : Tree α) : Prop :=
+  Traversable.foldl
 
 end Invariants
 
 
 /-! ### Accessor Lemmas for ForallTree -/
 section ForallTreeAccessors
 
-@[simp] lemma ForallTree.left_sub {p : α → Prop} {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma ForallTree.left_sub {p : α → Prop} {l : Tree α} {k : α} {r : Tree α}
     (h : ForallTree p (.node l k r)) : ForallTree p l := by
   cases h with | node _ _ _ hl _ _ => exact hl
 
-@[simp] lemma ForallTree.root {p : α → Prop} {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma ForallTree.root {p : α → Prop} {l : Tree α} {k : α} {r : Tree α}
     (h : ForallTree p (.node l k r)) : p k := by
   cases h with | node _ _ _ _ hk _ => exact hk
 
-@[simp] lemma ForallTree.right_sub {p : α → Prop} {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma ForallTree.right_sub {p : α → Prop} {l : Tree α} {k : α} {r : Tree α}
     (h : ForallTree p (.node l k r)) : ForallTree p r := by
   cases h with | node _ _ _ _ _ hr => exact hr
 
@@ -219,40 +182,25 @@ end ForallTreeAccessors
 /-! ### Accessor Lemmas for IsBST -/
 section IsBSTAccessors
 
-@[simp] lemma IsBST.forallTree_left [LinearOrder α] {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma IsBST.forallTree_left [LinearOrder α] {l : Tree α} {k : α} {r : Tree α}
     (h : IsBST (.node l k r)) : ForallTree (· < k) l := by
   cases h with | node _ _ _ hl _ _ _ => exact hl
 
-@[simp] lemma IsBST.forallTree_right [LinearOrder α] {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma IsBST.forallTree_right [LinearOrder α] {l : Tree α} {k : α} {r : Tree α}
     (h : IsBST (.node l k r)) : ForallTree (k < ·) r := by
   cases h with | node _ _ _ _ hr _ _ => exact hr
 
-@[simp] lemma IsBST.left_bst [LinearOrder α] {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma IsBST.left_bst [LinearOrder α] {l : Tree α} {k : α} {r : Tree α}
     (h : IsBST (.node l k r)) : IsBST l := by
   cases h with | node _ _ _ _ _ hl _ => exact hl
 
-@[simp] lemma IsBST.right_bst [LinearOrder α] {l : BinaryTree α} {k : α} {r : BinaryTree α}
+@[simp] lemma IsBST.right_bst [LinearOrder α] {l : Tree α} {k : α} {r : Tree α}
     (h : IsBST (.node l k r)) : IsBST r := by
   cases h with | node _ _ _ _ _ _ hr => exact hr
 
 end IsBSTAccessors
 
-end BinaryTree
-
-
-/-! ### BST Structure -/
-section BSTStructure
-
-structure BST (α : Type) [LinearOrder α] where
-  tree : BinaryTree α
-  hBST : BinaryTree.IsBST tree
-
-namespace BST
-
-/-- Checks if the BST contains a given key by delegating to the underlying tree. -/
-def contains [LinearOrder α] (t : BST α) (q : α) : Prop :=
-  t.tree.contains q
 
-end BST
+end Tree
 
-end BSTStructure
+end Cslib

Cslib/Foundations/Data/SplayTree/BSTAPI.lean

@@ -5,8 +5,10 @@ Authors: Anton Kovsharov, Antoine du Fresne von Hohenesche,
   Sorrachai Yingchareonthawornchai
 -/
 
-import Cslib.Foundations.Data.SplayTree.Complexity
-import Cslib.Foundations.Data.SplayTree.Correctness
+module
+
+public import Cslib.Foundations.Data.SplayTree.Complexity
+public import Cslib.Foundations.Data.SplayTree.Correctness
 
 /-!
 # Splay Tree API for BST
@@ -17,13 +19,17 @@ proofs, allowing users to safely and cleanly manipulate BSTs without manually
 re-proving the `IsBST` invariant after every rotation.
 -/
 
-variable {α : Type} [LinearOrder α]
+@[expose] public section
+
+namespace Cslib
 
 namespace SplayTree.BSTAPI
 
 open SplayTree
 open BinaryTree
 
+variable {α : Type} [LinearOrder α]
+
 /-! ### Core Operation -/
 
 /--
@@ -85,3 +91,5 @@ theorem nlogn_cost (n m : ℕ) (X : Fin m → α)
   SplayTree.nlogn_cost n m X init.tree h_size
 
 end SplayTree.BSTAPI
+
+end Cslib