...

Source file src/container/heap/heap.go

Documentation: container/heap

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

View as plain text