1 /*-------------------------------------------------------------------------
2 *
3 * binaryheap.c
4 * A simple binary heap implementation
5 *
6 * Portions Copyright (c) 2012-2016, PostgreSQL Global Development Group
7 *
8 * IDENTIFICATION
9 * src/backend/lib/binaryheap.c
10 *
11 *-------------------------------------------------------------------------
12 */
13
14 #include "postgres.h"
15
16 #include <math.h>
17
18 #include "lib/binaryheap.h"
19
20 static void sift_down(binaryheap *heap, int node_off);
21 static void sift_up(binaryheap *heap, int node_off);
22 static inline void swap_nodes(binaryheap *heap, int a, int b);
23
24 /*
25 * binaryheap_allocate
26 *
27 * Returns a pointer to a newly-allocated heap that has the capacity to
28 * store the given number of nodes, with the heap property defined by
29 * the given comparator function, which will be invoked with the additional
30 * argument specified by 'arg'.
31 */
32 binaryheap *
binaryheap_allocate(int capacity,binaryheap_comparator compare,void * arg)33 binaryheap_allocate(int capacity, binaryheap_comparator compare, void *arg)
34 {
35 int sz;
36 binaryheap *heap;
37
38 sz = offsetof(binaryheap, bh_nodes) +sizeof(Datum) * capacity;
39 heap = (binaryheap *) palloc(sz);
40 heap->bh_space = capacity;
41 heap->bh_compare = compare;
42 heap->bh_arg = arg;
43
44 heap->bh_size = 0;
45 heap->bh_has_heap_property = true;
46
47 return heap;
48 }
49
50 /*
51 * binaryheap_reset
52 *
53 * Resets the heap to an empty state, losing its data content but not the
54 * parameters passed at allocation.
55 */
56 void
binaryheap_reset(binaryheap * heap)57 binaryheap_reset(binaryheap *heap)
58 {
59 heap->bh_size = 0;
60 heap->bh_has_heap_property = true;
61 }
62
63 /*
64 * binaryheap_free
65 *
66 * Releases memory used by the given binaryheap.
67 */
68 void
binaryheap_free(binaryheap * heap)69 binaryheap_free(binaryheap *heap)
70 {
71 pfree(heap);
72 }
73
74 /*
75 * These utility functions return the offset of the left child, right
76 * child, and parent of the node at the given index, respectively.
77 *
78 * The heap is represented as an array of nodes, with the root node
79 * stored at index 0. The left child of node i is at index 2*i+1, and
80 * the right child at 2*i+2. The parent of node i is at index (i-1)/2.
81 */
82
83 static inline int
left_offset(int i)84 left_offset(int i)
85 {
86 return 2 * i + 1;
87 }
88
89 static inline int
right_offset(int i)90 right_offset(int i)
91 {
92 return 2 * i + 2;
93 }
94
95 static inline int
parent_offset(int i)96 parent_offset(int i)
97 {
98 return (i - 1) / 2;
99 }
100
101 /*
102 * binaryheap_add_unordered
103 *
104 * Adds the given datum to the end of the heap's list of nodes in O(1) without
105 * preserving the heap property. This is a convenience to add elements quickly
106 * to a new heap. To obtain a valid heap, one must call binaryheap_build()
107 * afterwards.
108 */
109 void
binaryheap_add_unordered(binaryheap * heap,Datum d)110 binaryheap_add_unordered(binaryheap *heap, Datum d)
111 {
112 if (heap->bh_size >= heap->bh_space)
113 elog(ERROR, "out of binary heap slots");
114 heap->bh_has_heap_property = false;
115 heap->bh_nodes[heap->bh_size] = d;
116 heap->bh_size++;
117 }
118
119 /*
120 * binaryheap_build
121 *
122 * Assembles a valid heap in O(n) from the nodes added by
123 * binaryheap_add_unordered(). Not needed otherwise.
124 */
125 void
binaryheap_build(binaryheap * heap)126 binaryheap_build(binaryheap *heap)
127 {
128 int i;
129
130 for (i = parent_offset(heap->bh_size - 1); i >= 0; i--)
131 sift_down(heap, i);
132 heap->bh_has_heap_property = true;
133 }
134
135 /*
136 * binaryheap_add
137 *
138 * Adds the given datum to the heap in O(log n) time, while preserving
139 * the heap property.
140 */
141 void
binaryheap_add(binaryheap * heap,Datum d)142 binaryheap_add(binaryheap *heap, Datum d)
143 {
144 if (heap->bh_size >= heap->bh_space)
145 elog(ERROR, "out of binary heap slots");
146 heap->bh_nodes[heap->bh_size] = d;
147 heap->bh_size++;
148 sift_up(heap, heap->bh_size - 1);
149 }
150
151 /*
152 * binaryheap_first
153 *
154 * Returns a pointer to the first (root, topmost) node in the heap
155 * without modifying the heap. The caller must ensure that this
156 * routine is not used on an empty heap. Always O(1).
157 */
158 Datum
binaryheap_first(binaryheap * heap)159 binaryheap_first(binaryheap *heap)
160 {
161 Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
162 return heap->bh_nodes[0];
163 }
164
165 /*
166 * binaryheap_remove_first
167 *
168 * Removes the first (root, topmost) node in the heap and returns a
169 * pointer to it after rebalancing the heap. The caller must ensure
170 * that this routine is not used on an empty heap. O(log n) worst
171 * case.
172 */
173 Datum
binaryheap_remove_first(binaryheap * heap)174 binaryheap_remove_first(binaryheap *heap)
175 {
176 Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
177
178 if (heap->bh_size == 1)
179 {
180 heap->bh_size--;
181 return heap->bh_nodes[0];
182 }
183
184 /*
185 * Swap the root and last nodes, decrease the size of the heap (i.e.
186 * remove the former root node) and sift the new root node down to its
187 * correct position.
188 */
189 swap_nodes(heap, 0, heap->bh_size - 1);
190 heap->bh_size--;
191 sift_down(heap, 0);
192
193 return heap->bh_nodes[heap->bh_size];
194 }
195
196 /*
197 * binaryheap_replace_first
198 *
199 * Replace the topmost element of a non-empty heap, preserving the heap
200 * property. O(1) in the best case, or O(log n) if it must fall back to
201 * sifting the new node down.
202 */
203 void
binaryheap_replace_first(binaryheap * heap,Datum d)204 binaryheap_replace_first(binaryheap *heap, Datum d)
205 {
206 Assert(!binaryheap_empty(heap) && heap->bh_has_heap_property);
207
208 heap->bh_nodes[0] = d;
209
210 if (heap->bh_size > 1)
211 sift_down(heap, 0);
212 }
213
214 /*
215 * Swap the contents of two nodes.
216 */
217 static inline void
swap_nodes(binaryheap * heap,int a,int b)218 swap_nodes(binaryheap *heap, int a, int b)
219 {
220 Datum swap;
221
222 swap = heap->bh_nodes[a];
223 heap->bh_nodes[a] = heap->bh_nodes[b];
224 heap->bh_nodes[b] = swap;
225 }
226
227 /*
228 * Sift a node up to the highest position it can hold according to the
229 * comparator.
230 */
231 static void
sift_up(binaryheap * heap,int node_off)232 sift_up(binaryheap *heap, int node_off)
233 {
234 while (node_off != 0)
235 {
236 int cmp;
237 int parent_off;
238
239 /*
240 * If this node is smaller than its parent, the heap condition is
241 * satisfied, and we're done.
242 */
243 parent_off = parent_offset(node_off);
244 cmp = heap->bh_compare(heap->bh_nodes[node_off],
245 heap->bh_nodes[parent_off],
246 heap->bh_arg);
247 if (cmp <= 0)
248 break;
249
250 /*
251 * Otherwise, swap the node and its parent and go on to check the
252 * node's new parent.
253 */
254 swap_nodes(heap, node_off, parent_off);
255 node_off = parent_off;
256 }
257 }
258
259 /*
260 * Sift a node down from its current position to satisfy the heap
261 * property.
262 */
263 static void
sift_down(binaryheap * heap,int node_off)264 sift_down(binaryheap *heap, int node_off)
265 {
266 while (true)
267 {
268 int left_off = left_offset(node_off);
269 int right_off = right_offset(node_off);
270 int swap_off = 0;
271
272 /* Is the left child larger than the parent? */
273 if (left_off < heap->bh_size &&
274 heap->bh_compare(heap->bh_nodes[node_off],
275 heap->bh_nodes[left_off],
276 heap->bh_arg) < 0)
277 swap_off = left_off;
278
279 /* Is the right child larger than the parent? */
280 if (right_off < heap->bh_size &&
281 heap->bh_compare(heap->bh_nodes[node_off],
282 heap->bh_nodes[right_off],
283 heap->bh_arg) < 0)
284 {
285 /* swap with the larger child */
286 if (!swap_off ||
287 heap->bh_compare(heap->bh_nodes[left_off],
288 heap->bh_nodes[right_off],
289 heap->bh_arg) < 0)
290 swap_off = right_off;
291 }
292
293 /*
294 * If we didn't find anything to swap, the heap condition is
295 * satisfied, and we're done.
296 */
297 if (!swap_off)
298 break;
299
300 /*
301 * Otherwise, swap the node with the child that violates the heap
302 * property; then go on to check its children.
303 */
304 swap_nodes(heap, swap_off, node_off);
305 node_off = swap_off;
306 }
307 }
308