148 lines
5.1 KiB
Text
148 lines
5.1 KiB
Text
/-
|
||
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
prelude
|
||
import init.data.list init.coe
|
||
universes u
|
||
|
||
namespace BinomialHeapImp
|
||
|
||
structure HeapNodeAux (α : Type u) (h : Type u) :=
|
||
(val : α) (rank : Nat) (children : List h)
|
||
|
||
inductive Heap (α : Type u) : Type u
|
||
| empty {} : Heap
|
||
| heap (ns : List (HeapNodeAux α Heap)) : Heap
|
||
|
||
abbrev HeapNode (α) := HeapNodeAux α (Heap α)
|
||
|
||
variables {α : Type u}
|
||
|
||
instance : Inhabited (Heap α) := ⟨Heap.empty⟩
|
||
|
||
def hRank : List (HeapNode α) → Nat
|
||
| [] => 0
|
||
| h::_ => h.rank
|
||
|
||
def isEmpty : Heap α → Bool
|
||
| Heap.empty => true
|
||
| _ => false
|
||
|
||
def singleton (a : α) : Heap α :=
|
||
Heap.heap [{ val := a, rank := 1, children := [] }]
|
||
|
||
@[specialize] def combine (lt : α → α → Bool) (n₁ n₂ : HeapNode α) : HeapNode α :=
|
||
if lt n₂.val n₁.val then
|
||
{ rank := n₂.rank + 1, children := n₂.children ++ [Heap.heap [n₁]], .. n₂ }
|
||
else
|
||
{ rank := n₁.rank + 1, children := n₁.children ++ [Heap.heap [n₂]], .. n₁ }
|
||
|
||
@[specialize] partial def mergeNodes (lt : α → α → Bool) : List (HeapNode α) → List (HeapNode α) → List (HeapNode α)
|
||
| [], h => h
|
||
| h, [] => h
|
||
| f@(h₁ :: t₁), s@(h₂ :: t₂) =>
|
||
if h₁.rank < h₂.rank then h₁ :: mergeNodes t₁ s
|
||
else if h₂.rank < h₁.rank then h₂ :: mergeNodes t₂ f
|
||
else
|
||
let merged := combine lt h₁ h₂;
|
||
let r := merged.rank;
|
||
if r != hRank t₁ then
|
||
if r != hRank t₂ then merged :: mergeNodes t₁ t₂ else mergeNodes (merged :: t₁) t₂
|
||
else
|
||
if r != hRank t₂ then mergeNodes t₁ (merged :: t₂) else merged :: mergeNodes t₁ t₂
|
||
|
||
@[specialize] def merge (lt : α → α → Bool) : Heap α → Heap α → Heap α
|
||
| Heap.empty, h => h
|
||
| h, Heap.empty => h
|
||
| Heap.heap h₁, Heap.heap h₂ => Heap.heap (mergeNodes lt h₁ h₂)
|
||
|
||
@[specialize] def headOpt (lt : α → α → Bool) : Heap α → Option α
|
||
| Heap.empty => none
|
||
| Heap.heap h => h.foldl
|
||
(fun r n => match r with
|
||
| none => n.val
|
||
| some v => if lt v n.val then v else n.val) none
|
||
|
||
/- O(log n) -/
|
||
@[specialize] def head [Inhabited α] (lt : α → α → Bool) : Heap α → α
|
||
| Heap.empty => default α
|
||
| Heap.heap [] => default α
|
||
| Heap.heap (h::hs) => hs.foldl (fun r n => if lt r n.val then r else n.val) h.val
|
||
|
||
@[specialize] def findMin (lt : α → α → Bool) : List (HeapNode α) → Nat → HeapNode α × Nat → HeapNode α × Nat
|
||
| [], _, r => r
|
||
| h::hs, idx, (h', idx') => if lt h.val h'.val then findMin hs (idx+1) (h, idx) else findMin hs (idx+1) (h', idx')
|
||
|
||
def tail (lt : α → α → Bool) : Heap α → Heap α
|
||
| Heap.empty => Heap.empty
|
||
| Heap.heap [] => Heap.empty
|
||
| Heap.heap [h] =>
|
||
match h.children with
|
||
| [] => Heap.empty
|
||
| (h::hs) => hs.foldl (merge lt) h
|
||
| Heap.heap hhs@(h::hs) =>
|
||
let (min, minIdx) := findMin lt hs 1 (h, 0);
|
||
let rest := hhs.eraseIdx minIdx;
|
||
min.children.foldl (merge lt) (Heap.heap rest)
|
||
|
||
partial def toList (lt : α → α → Bool) : Heap α → List α
|
||
| Heap.empty => []
|
||
| h => match headOpt lt h with
|
||
| none => []
|
||
| some a => a :: toList (tail lt h)
|
||
|
||
inductive WellFormed (lt : α → α → Bool) : Heap α → Prop
|
||
| emptyWff : WellFormed Heap.empty
|
||
| singletonWff (a : α) : WellFormed (singleton a)
|
||
| mergeWff (h₁ h₂ : Heap α) : WellFormed h₁ → WellFormed h₂ → WellFormed (merge lt h₁ h₂)
|
||
| tailWff (h : Heap α) : WellFormed h → WellFormed (tail lt h)
|
||
|
||
end BinomialHeapImp
|
||
|
||
open BinomialHeapImp
|
||
|
||
def BinomialHeap (α : Type u) (lt : α → α → Bool) := { h : Heap α // WellFormed lt h }
|
||
|
||
@[inline] def mkBinomialHeap (α : Type u) (lt : α → α → Bool) : BinomialHeap α lt :=
|
||
⟨Heap.empty, WellFormed.emptyWff lt⟩
|
||
|
||
namespace BinomialHeap
|
||
variables {α : Type u} {lt : α → α → Bool}
|
||
|
||
@[inline] def empty : BinomialHeap α lt :=
|
||
mkBinomialHeap α lt
|
||
|
||
@[inline] def isEmpty : BinomialHeap α lt → Bool
|
||
| ⟨b, _⟩ => BinomialHeapImp.isEmpty b
|
||
|
||
/- O(1) -/
|
||
@[inline] def singleton (a : α) : BinomialHeap α lt :=
|
||
⟨BinomialHeapImp.singleton a, WellFormed.singletonWff lt a⟩
|
||
|
||
/- O(log n) -/
|
||
@[inline] def merge : BinomialHeap α lt → BinomialHeap α lt → BinomialHeap α lt
|
||
| ⟨b₁, h₁⟩, ⟨b₂, h₂⟩ => ⟨BinomialHeapImp.merge lt b₁ b₂, WellFormed.mergeWff b₁ b₂ h₁ h₂⟩
|
||
|
||
/- O(log n) -/
|
||
@[inline] def head [Inhabited α] : BinomialHeap α lt → α
|
||
| ⟨b, _⟩ => BinomialHeapImp.head lt b
|
||
|
||
/- O(log n) -/
|
||
@[inline] def headOpt : BinomialHeap α lt → Option α
|
||
| ⟨b, _⟩ => BinomialHeapImp.headOpt lt b
|
||
|
||
/- O(log n) -/
|
||
@[inline] def tail : BinomialHeap α lt → BinomialHeap α lt
|
||
| ⟨b, h⟩ => ⟨BinomialHeapImp.tail lt b, WellFormed.tailWff b h⟩
|
||
|
||
/- O(log n) -/
|
||
@[inline] def insert (a : α) (h : BinomialHeap α lt) : BinomialHeap α lt :=
|
||
merge (singleton a) h
|
||
|
||
/- O(n log n) -/
|
||
@[inline] def toList : BinomialHeap α lt → List α
|
||
| ⟨b, _⟩ => BinomialHeapImp.toList lt b
|
||
|
||
end BinomialHeap
|