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
  | 
/-
  This file contains the Binary Heap type and its basic operations.
-/
import BinaryHeap.CompleteTree
structure BinaryHeap (α : Type u) (le : α → α → Bool) (n : Nat) where
  tree : BinaryHeap.CompleteTree α n
  valid : BinaryHeap.HeapPredicate tree le
  wellDefinedLe : BinaryHeap.transitive_le le ∧ BinaryHeap.total_le le
deriving Repr
namespace BinaryHeap
/-- Creates an empty heap. O(1) -/
def new (α : Type u) {le : α → α → Bool} (h₁ : transitive_le le) (h₂ : total_le le) : BinaryHeap α le 0 :=
  {
    tree := CompleteTree.empty,
    valid := CompleteTree.emptyIsHeap le
    wellDefinedLe := ⟨h₁,h₂⟩
  }
/-- Adds an element into a heap. O(log(n)) -/
def push {α : Type u} {le : α → α → Bool} {n : Nat} : α → BinaryHeap α le n → BinaryHeap α le (n+1)
| elem, .mk tree valid wellDefinedLe =>
  let valid := tree.heapPushIsHeap valid wellDefinedLe.left wellDefinedLe.right
  let tree := tree.heapPush le elem
  {tree, valid, wellDefinedLe}
/-- Removes the smallest element from the heap and returns it. O(log(n)) -/
def pop {α : Type u} {le : α → α → Bool} {n : Nat} : (BinaryHeap α le (n+1)) → (BinaryHeap α le n × α)
| {tree, valid, wellDefinedLe} =>
  let result := tree.heapPop le
  let resultValid := CompleteTree.heapPopIsHeap le tree valid wellDefinedLe
  ({ tree := result.fst, valid := resultValid, wellDefinedLe}, result.snd)
/-- Removes the element at the given (depth-first) index from the heap and returns it. O(log(n)) -/
def removeAt {α : Type u} {le : α → α → Bool} {n : Nat} : (BinaryHeap α le (n+1)) → (Fin (n+1)) → (BinaryHeap α le n × α)
| {tree, valid, wellDefinedLe}, index =>
  let result := tree.heapRemoveAt le index
  let resultValid := CompleteTree.heapRemoveAtIsHeap le index tree valid wellDefinedLe
  ({ tree := result.fst, valid := resultValid, wellDefinedLe}, result.snd)
/-- Updates the element at the given (depth-first) index and returns its previous value. O(log(n)) -/
def updateAt {α : Type u} {le : α → α → Bool} {n : Nat} : (BinaryHeap α le n) → (Fin n) → α → (BinaryHeap α le n × α)
| {tree, valid, wellDefinedLe}, index, value =>
  let result := tree.heapUpdateAt le index value
  let resultValid := CompleteTree.heapUpdateAtIsHeap le index value tree valid wellDefinedLe.left wellDefinedLe.right
  ({ tree := result.fst, valid := resultValid, wellDefinedLe}, result.snd)
/-- Searches for an element in the heap and returns its index. **O(n)** -/
def indexOf {α : Type u} {le : α → α → Bool} {n : Nat} : BinaryHeap α le n → (pred : α → Bool) → Option (Fin n)
| {tree, valid := _, wellDefinedLe := _}, pred =>
  tree.indexOf pred
instance : GetElem (BinaryHeap α le n) (Nat) α (λ _ index ↦ index < n) where
  getElem := λ heap index h₁ ↦ heap.tree.get ⟨index, h₁⟩
instance : GetElem (BinaryHeap α le n) (Fin n) α (λ _ _ ↦ True) where
  getElem := λ heap index _ ↦ heap.tree.get index
/--Helper for the common case of using natural numbers for sorting.-/
theorem nat_ble_to_heap_le_total : total_le Nat.ble :=
  λa b ↦ Nat.ble_eq.substr $ Nat.ble_eq.substr (Nat.le_total a b)
/--Helper for the common case of using natural numbers for sorting.-/
theorem nat_ble_to_heap_transitive_le : transitive_le Nat.ble :=
  λ a b c ↦
    Nat.le_trans (n := a) (m := b) (k := c)
    |> Nat.ble_eq.substr
    |> Nat.ble_eq.substr
    |> Nat.ble_eq.substr
    |> and_imp.mpr
 
  |