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//
17package heap
18
19import "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.
32type 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().
42func 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().
52func 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).
60func 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().
69func 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().
84func Fix(h Interface, i int) {
85	if !down(h, i, h.Len()) {
86		up(h, i)
87	}
88}
89
90func 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
101func 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