• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# -*- coding: latin-1 -*-
2
3"""Heap queue algorithm (a.k.a. priority queue).
4
5Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for
6all k, counting elements from 0.  For the sake of comparison,
7non-existing elements are considered to be infinite.  The interesting
8property of a heap is that a[0] is always its smallest element.
9
10Usage:
11
12heap = []            # creates an empty heap
13heappush(heap, item) # pushes a new item on the heap
14item = heappop(heap) # pops the smallest item from the heap
15item = heap[0]       # smallest item on the heap without popping it
16heapify(x)           # transforms list into a heap, in-place, in linear time
17item = heapreplace(heap, item) # pops and returns smallest item, and adds
18                               # new item; the heap size is unchanged
19
20Our API differs from textbook heap algorithms as follows:
21
22- We use 0-based indexing.  This makes the relationship between the
23  index for a node and the indexes for its children slightly less
24  obvious, but is more suitable since Python uses 0-based indexing.
25
26- Our heappop() method returns the smallest item, not the largest.
27
28These two make it possible to view the heap as a regular Python list
29without surprises: heap[0] is the smallest item, and heap.sort()
30maintains the heap invariant!
31"""
32
33# Original code by Kevin O'Connor, augmented by Tim Peters and Raymond Hettinger
34
35__about__ = """Heap queues
36
37[explanation by Fran�ois Pinard]
38
39Heaps are arrays for which a[k] <= a[2*k+1] and a[k] <= a[2*k+2] for
40all k, counting elements from 0.  For the sake of comparison,
41non-existing elements are considered to be infinite.  The interesting
42property of a heap is that a[0] is always its smallest element.
43
44The strange invariant above is meant to be an efficient memory
45representation for a tournament.  The numbers below are `k', not a[k]:
46
47                                   0
48
49                  1                                 2
50
51          3               4                5               6
52
53      7       8       9       10      11      12      13      14
54
55    15 16   17 18   19 20   21 22   23 24   25 26   27 28   29 30
56
57
58In the tree above, each cell `k' is topping `2*k+1' and `2*k+2'.  In
59a usual binary tournament we see in sports, each cell is the winner
60over the two cells it tops, and we can trace the winner down the tree
61to see all opponents s/he had.  However, in many computer applications
62of such tournaments, we do not need to trace the history of a winner.
63To be more memory efficient, when a winner is promoted, we try to
64replace it by something else at a lower level, and the rule becomes
65that a cell and the two cells it tops contain three different items,
66but the top cell "wins" over the two topped cells.
67
68If this heap invariant is protected at all time, index 0 is clearly
69the overall winner.  The simplest algorithmic way to remove it and
70find the "next" winner is to move some loser (let's say cell 30 in the
71diagram above) into the 0 position, and then percolate this new 0 down
72the tree, exchanging values, until the invariant is re-established.
73This is clearly logarithmic on the total number of items in the tree.
74By iterating over all items, you get an O(n ln n) sort.
75
76A nice feature of this sort is that you can efficiently insert new
77items while the sort is going on, provided that the inserted items are
78not "better" than the last 0'th element you extracted.  This is
79especially useful in simulation contexts, where the tree holds all
80incoming events, and the "win" condition means the smallest scheduled
81time.  When an event schedule other events for execution, they are
82scheduled into the future, so they can easily go into the heap.  So, a
83heap is a good structure for implementing schedulers (this is what I
84used for my MIDI sequencer :-).
85
86Various structures for implementing schedulers have been extensively
87studied, and heaps are good for this, as they are reasonably speedy,
88the speed is almost constant, and the worst case is not much different
89than the average case.  However, there are other representations which
90are more efficient overall, yet the worst cases might be terrible.
91
92Heaps are also very useful in big disk sorts.  You most probably all
93know that a big sort implies producing "runs" (which are pre-sorted
94sequences, which size is usually related to the amount of CPU memory),
95followed by a merging passes for these runs, which merging is often
96very cleverly organised[1].  It is very important that the initial
97sort produces the longest runs possible.  Tournaments are a good way
98to that.  If, using all the memory available to hold a tournament, you
99replace and percolate items that happen to fit the current run, you'll
100produce runs which are twice the size of the memory for random input,
101and much better for input fuzzily ordered.
102
103Moreover, if you output the 0'th item on disk and get an input which
104may not fit in the current tournament (because the value "wins" over
105the last output value), it cannot fit in the heap, so the size of the
106heap decreases.  The freed memory could be cleverly reused immediately
107for progressively building a second heap, which grows at exactly the
108same rate the first heap is melting.  When the first heap completely
109vanishes, you switch heaps and start a new run.  Clever and quite
110effective!
111
112In a word, heaps are useful memory structures to know.  I use them in
113a few applications, and I think it is good to keep a `heap' module
114around. :-)
115
116--------------------
117[1] The disk balancing algorithms which are current, nowadays, are
118more annoying than clever, and this is a consequence of the seeking
119capabilities of the disks.  On devices which cannot seek, like big
120tape drives, the story was quite different, and one had to be very
121clever to ensure (far in advance) that each tape movement will be the
122most effective possible (that is, will best participate at
123"progressing" the merge).  Some tapes were even able to read
124backwards, and this was also used to avoid the rewinding time.
125Believe me, real good tape sorts were quite spectacular to watch!
126From all times, sorting has always been a Great Art! :-)
127"""
128
129__all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge',
130           'nlargest', 'nsmallest', 'heappushpop']
131
132from itertools import islice, count, imap, izip, tee, chain
133from operator import itemgetter
134
135def cmp_lt(x, y):
136    # Use __lt__ if available; otherwise, try __le__.
137    # In Py3.x, only __lt__ will be called.
138    return (x < y) if hasattr(x, '__lt__') else (not y <= x)
139
140def heappush(heap, item):
141    """Push item onto heap, maintaining the heap invariant."""
142    heap.append(item)
143    _siftdown(heap, 0, len(heap)-1)
144
145def heappop(heap):
146    """Pop the smallest item off the heap, maintaining the heap invariant."""
147    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
148    if heap:
149        returnitem = heap[0]
150        heap[0] = lastelt
151        _siftup(heap, 0)
152    else:
153        returnitem = lastelt
154    return returnitem
155
156def heapreplace(heap, item):
157    """Pop and return the current smallest value, and add the new item.
158
159    This is more efficient than heappop() followed by heappush(), and can be
160    more appropriate when using a fixed-size heap.  Note that the value
161    returned may be larger than item!  That constrains reasonable uses of
162    this routine unless written as part of a conditional replacement:
163
164        if item > heap[0]:
165            item = heapreplace(heap, item)
166    """
167    returnitem = heap[0]    # raises appropriate IndexError if heap is empty
168    heap[0] = item
169    _siftup(heap, 0)
170    return returnitem
171
172def heappushpop(heap, item):
173    """Fast version of a heappush followed by a heappop."""
174    if heap and cmp_lt(heap[0], item):
175        item, heap[0] = heap[0], item
176        _siftup(heap, 0)
177    return item
178
179def heapify(x):
180    """Transform list into a heap, in-place, in O(len(x)) time."""
181    n = len(x)
182    # Transform bottom-up.  The largest index there's any point to looking at
183    # is the largest with a child index in-range, so must have 2*i + 1 < n,
184    # or i < (n-1)/2.  If n is even = 2*j, this is (2*j-1)/2 = j-1/2 so
185    # j-1 is the largest, which is n//2 - 1.  If n is odd = 2*j+1, this is
186    # (2*j+1-1)/2 = j so j-1 is the largest, and that's again n//2-1.
187    for i in reversed(xrange(n//2)):
188        _siftup(x, i)
189
190def _heappushpop_max(heap, item):
191    """Maxheap version of a heappush followed by a heappop."""
192    if heap and cmp_lt(item, heap[0]):
193        item, heap[0] = heap[0], item
194        _siftup_max(heap, 0)
195    return item
196
197def _heapify_max(x):
198    """Transform list into a maxheap, in-place, in O(len(x)) time."""
199    n = len(x)
200    for i in reversed(range(n//2)):
201        _siftup_max(x, i)
202
203def nlargest(n, iterable):
204    """Find the n largest elements in a dataset.
205
206    Equivalent to:  sorted(iterable, reverse=True)[:n]
207    """
208    if n < 0:
209        return []
210    it = iter(iterable)
211    result = list(islice(it, n))
212    if not result:
213        return result
214    heapify(result)
215    _heappushpop = heappushpop
216    for elem in it:
217        _heappushpop(result, elem)
218    result.sort(reverse=True)
219    return result
220
221def nsmallest(n, iterable):
222    """Find the n smallest elements in a dataset.
223
224    Equivalent to:  sorted(iterable)[:n]
225    """
226    if n < 0:
227        return []
228    it = iter(iterable)
229    result = list(islice(it, n))
230    if not result:
231        return result
232    _heapify_max(result)
233    _heappushpop = _heappushpop_max
234    for elem in it:
235        _heappushpop(result, elem)
236    result.sort()
237    return result
238
239# 'heap' is a heap at all indices >= startpos, except possibly for pos.  pos
240# is the index of a leaf with a possibly out-of-order value.  Restore the
241# heap invariant.
242def _siftdown(heap, startpos, pos):
243    newitem = heap[pos]
244    # Follow the path to the root, moving parents down until finding a place
245    # newitem fits.
246    while pos > startpos:
247        parentpos = (pos - 1) >> 1
248        parent = heap[parentpos]
249        if cmp_lt(newitem, parent):
250            heap[pos] = parent
251            pos = parentpos
252            continue
253        break
254    heap[pos] = newitem
255
256# The child indices of heap index pos are already heaps, and we want to make
257# a heap at index pos too.  We do this by bubbling the smaller child of
258# pos up (and so on with that child's children, etc) until hitting a leaf,
259# then using _siftdown to move the oddball originally at index pos into place.
260#
261# We *could* break out of the loop as soon as we find a pos where newitem <=
262# both its children, but turns out that's not a good idea, and despite that
263# many books write the algorithm that way.  During a heap pop, the last array
264# element is sifted in, and that tends to be large, so that comparing it
265# against values starting from the root usually doesn't pay (= usually doesn't
266# get us out of the loop early).  See Knuth, Volume 3, where this is
267# explained and quantified in an exercise.
268#
269# Cutting the # of comparisons is important, since these routines have no
270# way to extract "the priority" from an array element, so that intelligence
271# is likely to be hiding in custom __cmp__ methods, or in array elements
272# storing (priority, record) tuples.  Comparisons are thus potentially
273# expensive.
274#
275# On random arrays of length 1000, making this change cut the number of
276# comparisons made by heapify() a little, and those made by exhaustive
277# heappop() a lot, in accord with theory.  Here are typical results from 3
278# runs (3 just to demonstrate how small the variance is):
279#
280# Compares needed by heapify     Compares needed by 1000 heappops
281# --------------------------     --------------------------------
282# 1837 cut to 1663               14996 cut to 8680
283# 1855 cut to 1659               14966 cut to 8678
284# 1847 cut to 1660               15024 cut to 8703
285#
286# Building the heap by using heappush() 1000 times instead required
287# 2198, 2148, and 2219 compares:  heapify() is more efficient, when
288# you can use it.
289#
290# The total compares needed by list.sort() on the same lists were 8627,
291# 8627, and 8632 (this should be compared to the sum of heapify() and
292# heappop() compares):  list.sort() is (unsurprisingly!) more efficient
293# for sorting.
294
295def _siftup(heap, pos):
296    endpos = len(heap)
297    startpos = pos
298    newitem = heap[pos]
299    # Bubble up the smaller child until hitting a leaf.
300    childpos = 2*pos + 1    # leftmost child position
301    while childpos < endpos:
302        # Set childpos to index of smaller child.
303        rightpos = childpos + 1
304        if rightpos < endpos and not cmp_lt(heap[childpos], heap[rightpos]):
305            childpos = rightpos
306        # Move the smaller child up.
307        heap[pos] = heap[childpos]
308        pos = childpos
309        childpos = 2*pos + 1
310    # The leaf at pos is empty now.  Put newitem there, and bubble it up
311    # to its final resting place (by sifting its parents down).
312    heap[pos] = newitem
313    _siftdown(heap, startpos, pos)
314
315def _siftdown_max(heap, startpos, pos):
316    'Maxheap variant of _siftdown'
317    newitem = heap[pos]
318    # Follow the path to the root, moving parents down until finding a place
319    # newitem fits.
320    while pos > startpos:
321        parentpos = (pos - 1) >> 1
322        parent = heap[parentpos]
323        if cmp_lt(parent, newitem):
324            heap[pos] = parent
325            pos = parentpos
326            continue
327        break
328    heap[pos] = newitem
329
330def _siftup_max(heap, pos):
331    'Maxheap variant of _siftup'
332    endpos = len(heap)
333    startpos = pos
334    newitem = heap[pos]
335    # Bubble up the larger child until hitting a leaf.
336    childpos = 2*pos + 1    # leftmost child position
337    while childpos < endpos:
338        # Set childpos to index of larger child.
339        rightpos = childpos + 1
340        if rightpos < endpos and not cmp_lt(heap[rightpos], heap[childpos]):
341            childpos = rightpos
342        # Move the larger child up.
343        heap[pos] = heap[childpos]
344        pos = childpos
345        childpos = 2*pos + 1
346    # The leaf at pos is empty now.  Put newitem there, and bubble it up
347    # to its final resting place (by sifting its parents down).
348    heap[pos] = newitem
349    _siftdown_max(heap, startpos, pos)
350
351# If available, use C implementation
352try:
353    from _heapq import *
354except ImportError:
355    pass
356
357def merge(*iterables):
358    '''Merge multiple sorted inputs into a single sorted output.
359
360    Similar to sorted(itertools.chain(*iterables)) but returns a generator,
361    does not pull the data into memory all at once, and assumes that each of
362    the input streams is already sorted (smallest to largest).
363
364    >>> list(merge([1,3,5,7], [0,2,4,8], [5,10,15,20], [], [25]))
365    [0, 1, 2, 3, 4, 5, 5, 7, 8, 10, 15, 20, 25]
366
367    '''
368    _heappop, _heapreplace, _StopIteration = heappop, heapreplace, StopIteration
369    _len = len
370
371    h = []
372    h_append = h.append
373    for itnum, it in enumerate(map(iter, iterables)):
374        try:
375            next = it.next
376            h_append([next(), itnum, next])
377        except _StopIteration:
378            pass
379    heapify(h)
380
381    while _len(h) > 1:
382        try:
383            while 1:
384                v, itnum, next = s = h[0]
385                yield v
386                s[0] = next()               # raises StopIteration when exhausted
387                _heapreplace(h, s)          # restore heap condition
388        except _StopIteration:
389            _heappop(h)                     # remove empty iterator
390    if h:
391        # fast case when only a single iterator remains
392        v, itnum, next = h[0]
393        yield v
394        for v in next.__self__:
395            yield v
396
397# Extend the implementations of nsmallest and nlargest to use a key= argument
398_nsmallest = nsmallest
399def nsmallest(n, iterable, key=None):
400    """Find the n smallest elements in a dataset.
401
402    Equivalent to:  sorted(iterable, key=key)[:n]
403    """
404    # Short-cut for n==1 is to use min() when len(iterable)>0
405    if n == 1:
406        it = iter(iterable)
407        head = list(islice(it, 1))
408        if not head:
409            return []
410        if key is None:
411            return [min(chain(head, it))]
412        return [min(chain(head, it), key=key)]
413
414    # When n>=size, it's faster to use sorted()
415    try:
416        size = len(iterable)
417    except (TypeError, AttributeError):
418        pass
419    else:
420        if n >= size:
421            return sorted(iterable, key=key)[:n]
422
423    # When key is none, use simpler decoration
424    if key is None:
425        it = izip(iterable, count())                        # decorate
426        result = _nsmallest(n, it)
427        return map(itemgetter(0), result)                   # undecorate
428
429    # General case, slowest method
430    in1, in2 = tee(iterable)
431    it = izip(imap(key, in1), count(), in2)                 # decorate
432    result = _nsmallest(n, it)
433    return map(itemgetter(2), result)                       # undecorate
434
435_nlargest = nlargest
436def nlargest(n, iterable, key=None):
437    """Find the n largest elements in a dataset.
438
439    Equivalent to:  sorted(iterable, key=key, reverse=True)[:n]
440    """
441
442    # Short-cut for n==1 is to use max() when len(iterable)>0
443    if n == 1:
444        it = iter(iterable)
445        head = list(islice(it, 1))
446        if not head:
447            return []
448        if key is None:
449            return [max(chain(head, it))]
450        return [max(chain(head, it), key=key)]
451
452    # When n>=size, it's faster to use sorted()
453    try:
454        size = len(iterable)
455    except (TypeError, AttributeError):
456        pass
457    else:
458        if n >= size:
459            return sorted(iterable, key=key, reverse=True)[:n]
460
461    # When key is none, use simpler decoration
462    if key is None:
463        it = izip(iterable, count(0,-1))                    # decorate
464        result = _nlargest(n, it)
465        return map(itemgetter(0), result)                   # undecorate
466
467    # General case, slowest method
468    in1, in2 = tee(iterable)
469    it = izip(imap(key, in1), count(0,-1), in2)             # decorate
470    result = _nlargest(n, it)
471    return map(itemgetter(2), result)                       # undecorate
472
473if __name__ == "__main__":
474    # Simple sanity test
475    heap = []
476    data = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
477    for item in data:
478        heappush(heap, item)
479    sort = []
480    while heap:
481        sort.append(heappop(heap))
482    print sort
483
484    import doctest
485    doctest.testmod()
486