feat(library/init/data): add BinomialHeap

library/init/data/binomialheap/basic.lean

@@ -0,0 +1,148 @@
+/-
+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

library/init/data/binomialheap/default.lean

@@ -0,0 +1,7 @@
+/-
+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.binomialheap.basic

tests/compiler/binomial.lean

@@ -0,0 +1,27 @@
+import init.data.binomialheap
+
+abbrev Heap := BinomialHeap Nat (fun a b => a < b)
+
+def mkHeap (n m : Nat) : Heap :=
+let hs : List Heap := n.fold (fun i hs =>
+  let h : Heap := BinomialHeap.empty;
+  let h : Heap := m.fold (fun j h =>
+    let v := n*m - j*n - i;
+    h.insert v) h;
+  h :: hs)
+  [];
+hs.foldl (fun h₁ h₂ => h₁.merge h₂) BinomialHeap.empty
+
+partial def display : Nat → Heap → IO Unit
+| prev h :=
+  if h.isEmpty then pure ()
+  else do
+    let m := h.head;
+    unless (prev < m) (IO.println ("failed " ++ toString prev ++ " " ++ toString m));
+    display m h.tail
+
+def main : IO Unit :=
+do
+let h := mkHeap 100 1000;
+display 0 h;
+IO.println h.toList.length

tests/compiler/binomial.lean.expected.out

@@ -0,0 +1 @@
+100000