heap.go 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  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. // Package heap provides heap operations for any type that implements
  5. // heap.Interface. A heap is a tree with the property that each node is the
  6. // minimum-valued node in its subtree.
  7. //
  8. // The minimum element in the tree is the root, at index 0.
  9. //
  10. // A heap is a common way to implement a priority queue. To build a priority
  11. // queue, implement the Heap interface with the (negative) priority as the
  12. // ordering for the Less method, so Push adds items while Pop removes the
  13. // highest-priority item from the queue. The Examples include such an
  14. // implementation; the file example_pq_test.go has the complete source.
  15. package heap
  16. import "sort"
  17. // The Interface type describes the requirements
  18. // for a type using the routines in this package.
  19. // Any type that implements it may be used as a
  20. // min-heap with the following invariants (established after
  21. // Init has been called or if the data is empty or sorted):
  22. //
  23. // !h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
  24. //
  25. // Note that Push and Pop in this interface are for package heap's
  26. // implementation to call. To add and remove things from the heap,
  27. // use heap.Push and heap.Pop.
  28. type Interface[T any] interface {
  29. sort.Interface
  30. Push(x T) // add x as element Len()
  31. Pop() T // remove and return element Len() - 1.
  32. }
  33. // Init establishes the heap invariants required by the other routines in this package.
  34. // Init is idempotent with respect to the heap invariants
  35. // and may be called whenever the heap invariants may have been invalidated.
  36. // The complexity is O(n) where n = h.Len().
  37. func Init[T any](h Interface[T]) {
  38. // heapify
  39. n := h.Len()
  40. for i := n/2 - 1; i >= 0; i-- {
  41. down(h, i, n)
  42. }
  43. }
  44. // Push pushes the element x onto the heap.
  45. // The complexity is O(log n) where n = h.Len().
  46. func Push[T any](h Interface[T], x T) {
  47. h.Push(x)
  48. up(h, h.Len()-1)
  49. }
  50. // Pop removes and returns the minimum element (according to Less) from the heap.
  51. // The complexity is O(log n) where n = h.Len().
  52. // Pop is equivalent to Remove(h, 0).
  53. func Pop[T any](h Interface[T]) T {
  54. n := h.Len() - 1
  55. h.Swap(0, n)
  56. down(h, 0, n)
  57. return h.Pop()
  58. }
  59. // Remove removes and returns the element at index i from the heap.
  60. // The complexity is O(log n) where n = h.Len().
  61. func Remove[T any](h Interface[T], i int) T {
  62. n := h.Len() - 1
  63. if n != i {
  64. h.Swap(i, n)
  65. if !down(h, i, n) {
  66. up(h, i)
  67. }
  68. }
  69. return h.Pop()
  70. }
  71. // Fix re-establishes the heap ordering after the element at index i has changed its value.
  72. // Changing the value of the element at index i and then calling Fix is equivalent to,
  73. // but less expensive than, calling Remove(h, i) followed by a Push of the new value.
  74. // The complexity is O(log n) where n = h.Len().
  75. func Fix[T any](h Interface[T], i int) {
  76. if !down(h, i, h.Len()) {
  77. up(h, i)
  78. }
  79. }
  80. func up[T any](h Interface[T], j int) {
  81. for {
  82. i := (j - 1) / 2 // parent
  83. if i == j || !h.Less(j, i) {
  84. break
  85. }
  86. h.Swap(i, j)
  87. j = i
  88. }
  89. }
  90. func down[T any](h Interface[T], i0, n int) bool {
  91. i := i0
  92. for {
  93. j1 := 2*i + 1
  94. if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
  95. break
  96. }
  97. j := j1 // left child
  98. if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
  99. j = j2 // = 2*i + 2 // right child
  100. }
  101. if !h.Less(j, i) {
  102. break
  103. }
  104. h.Swap(i, j)
  105. i = j
  106. }
  107. return i > i0
  108. }