diff --git a/library/init/data/binomialheap/basic.lean b/library/init/data/binomialheap/basic.lean new file mode 100644 index 0000000000..0441189c19 --- /dev/null +++ b/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 diff --git a/library/init/data/binomialheap/default.lean b/library/init/data/binomialheap/default.lean new file mode 100644 index 0000000000..a1af238505 --- /dev/null +++ b/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 diff --git a/tests/compiler/binomial.lean b/tests/compiler/binomial.lean new file mode 100644 index 0000000000..6d3fe7c8d0 --- /dev/null +++ b/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 diff --git a/tests/compiler/binomial.lean.expected.out b/tests/compiler/binomial.lean.expected.out new file mode 100644 index 0000000000..f7393e847d --- /dev/null +++ b/tests/compiler/binomial.lean.expected.out @@ -0,0 +1 @@ +100000