Source file src/pkg/container/heap/heap.go
1 // Copyright 2009 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 // Package heap provides heap operations for any type that implements 6 // heap.Interface. 7 // 8 package heap 9 10 import "sort" 11 12 // Any type that implements heap.Interface may be used as a 13 // min-heap with the following invariants (established after 14 // Init has been called): 15 // 16 // !h.Less(j, i) for 0 <= i < h.Len() and j = 2*i+1 or 2*i+2 and j < h.Len() 17 // 18 type Interface interface { 19 sort.Interface 20 Push(x interface{}) 21 Pop() interface{} 22 } 23 24 // A heap must be initialized before any of the heap operations 25 // can be used. Init is idempotent with respect to the heap invariants 26 // and may be called whenever the heap invariants may have been invalidated. 27 // Its complexity is O(n) where n = h.Len(). 28 // 29 func Init(h Interface) { 30 // heapify 31 n := h.Len() 32 for i := n/2 - 1; i >= 0; i-- { 33 down(h, i, n) 34 } 35 } 36 37 // Push pushes the element x onto the heap. The complexity is 38 // O(log(n)) where n = h.Len(). 39 // 40 func Push(h Interface, x interface{}) { 41 h.Push(x) 42 up(h, h.Len()-1) 43 } 44 45 // Pop removes the minimum element (according to Less) from the heap 46 // and returns it. The complexity is O(log(n)) where n = h.Len(). 47 // Same as Remove(h, 0). 48 // 49 func Pop(h Interface) interface{} { 50 n := h.Len() - 1 51 h.Swap(0, n) 52 down(h, 0, n) 53 return h.Pop() 54 } 55 56 // Remove removes the element at index i from the heap. 57 // The complexity is O(log(n)) where n = h.Len(). 58 // 59 func Remove(h Interface, i int) interface{} { 60 n := h.Len() - 1 61 if n != i { 62 h.Swap(i, n) 63 down(h, i, n) 64 up(h, i) 65 } 66 return h.Pop() 67 } 68 69 func up(h Interface, j int) { 70 for { 71 i := (j - 1) / 2 // parent 72 if i == j || h.Less(i, j) { 73 break 74 } 75 h.Swap(i, j) 76 j = i 77 } 78 } 79 80 func down(h Interface, i, n int) { 81 for { 82 j1 := 2*i + 1 83 if j1 >= n { 84 break 85 } 86 j := j1 // left child 87 if j2 := j1 + 1; j2 < n && !h.Less(j1, j2) { 88 j = j2 // = 2*i + 2 // right child 89 } 90 if h.Less(i, j) { 91 break 92 } 93 h.Swap(i, j) 94 i = j 95 } 96 }