diff --git a/Lib/heapq.py b/Lib/heapq.py --- a/Lib/heapq.py +++ b/Lib/heapq.py @@ -125,9 +125,17 @@ """ __all__ = ['heappush', 'heappop', 'heapify', 'heapreplace', 'merge', - 'nlargest', 'nsmallest', 'heappushpop'] + 'nlargest', 'nsmallest', 'heappushpop', 'Heap'] from itertools import islice, count, tee, chain +try: + from _thread import allocate_lock as Lock +except ImportError: + from _dummy_thread import allocate_lock as Lock + +################################################################################ +### MinHeap +################################################################################ def heappush(heap, item): """Push item onto heap, maintaining the heap invariant.""" @@ -179,55 +187,6 @@ for i in reversed(range(n//2)): _siftup(x, i) -def _heappushpop_max(heap, item): - """Maxheap version of a heappush followed by a heappop.""" - if heap and item < heap[0]: - item, heap[0] = heap[0], item - _siftup_max(heap, 0) - return item - -def _heapify_max(x): - """Transform list into a maxheap, in-place, in O(len(x)) time.""" - n = len(x) - for i in reversed(range(n//2)): - _siftup_max(x, i) - -def nlargest(n, iterable): - """Find the n largest elements in a dataset. - - Equivalent to: sorted(iterable, reverse=True)[:n] - """ - if n < 0: - return [] - it = iter(iterable) - result = list(islice(it, n)) - if not result: - return result - heapify(result) - _heappushpop = heappushpop - for elem in it: - _heappushpop(result, elem) - result.sort(reverse=True) - return result - -def nsmallest(n, iterable): - """Find the n smallest elements in a dataset. - - Equivalent to: sorted(iterable)[:n] - """ - if n < 0: - return [] - it = iter(iterable) - result = list(islice(it, n)) - if not result: - return result - _heapify_max(result) - _heappushpop = _heappushpop_max - for elem in it: - _heappushpop(result, elem) - result.sort() - return result - # 'heap' is a heap at all indices >= startpos, except possibly for pos. pos # is the index of a leaf with a possibly out-of-order value. Restore the # heap invariant. @@ -304,6 +263,41 @@ heap[pos] = newitem _siftdown(heap, startpos, pos) +################################################################################ +### MaxHeap +################################################################################ + +def _heappop_max(heap): + """Maxheap version of heappop""" + lastelt = heap.pop() # raises appropriate IndexError if heap is empty + if heap: + returnitem = heap[0] + heap[0] = lastelt + _siftup_max(heap, 0) + else: + returnitem = lastelt + return returnitem + +def _heapreplace_max(heap, item): + """Maxheap version of heapreplace""" + returnitem = heap[0] # raises appropriate IndexError if heap is empty + heap[0] = item + _siftup_max(heap, 0) + return returnitem + +def _heappushpop_max(heap, item): + """Maxheap version of a heappush followed by a heappop.""" + if heap and item < heap[0]: + item, heap[0] = heap[0], item + _siftup_max(heap, 0) + return item + +def _heapify_max(x): + """Transform list into a maxheap, in-place, in O(len(x)) time.""" + n = len(x) + for i in reversed(range(n//2)): + _siftup_max(x, i) + def _siftdown_max(heap, startpos, pos): 'Maxheap variant of _siftdown' newitem = heap[pos] @@ -340,12 +334,60 @@ heap[pos] = newitem _siftdown_max(heap, startpos, pos) -# If available, use C implementation +################################################################################ +### Basic nlargest() and nsmallest() without a key-function +################################################################################ + +def nlargest(n, iterable): + """Find the n largest elements in a dataset. + + Equivalent to: sorted(iterable, reverse=True)[:n] + """ + if n < 0: + return [] + it = iter(iterable) + result = list(islice(it, n)) + if not result: + return result + heapify(result) + _heappushpop = heappushpop + for elem in it: + _heappushpop(result, elem) + result.sort(reverse=True) + return result + +def nsmallest(n, iterable): + """Find the n smallest elements in a dataset. + + Equivalent to: sorted(iterable)[:n] + """ + if n < 0: + return [] + it = iter(iterable) + result = list(islice(it, n)) + if not result: + return result + _heapify_max(result) + _heappushpop = _heappushpop_max + for elem in it: + _heappushpop(result, elem) + result.sort() + return result + + +################################################################################ +### If available, use C implementation +################################################################################ + try: from _heapq import * except ImportError: pass +################################################################################ +### merge() +################################################################################ + def merge(*iterables): '''Merge multiple sorted inputs into a single sorted output. @@ -381,7 +423,10 @@ except IndexError: return -# Extend the implementations of nsmallest and nlargest to use a key= argument +################################################################################ +### Add key-function to nlargest() and nsmallest() +################################################################################ + _nsmallest = nsmallest def nsmallest(n, iterable, key=None): """Find the n smallest elements in a dataset. @@ -457,6 +502,81 @@ result = _nlargest(n, it) return [r[2] for r in result] # undecorate +################################################################################ +### Heap class +################################################################################ + +class Heap: + '''Fully encapsulated heap class + + Features: + * Thread-safe operations + * Comparison stability (equal keys returned in same order added) + * Key-functions (allowing the heap to work with non-comparable + objects such as functions) + * Minheap or maxheap + * Easily extendable using subclassing + + ''' + + def _wrap(self, value): + if self.key is None: + return value, self.count() + return self.key(value), self.count(), value + + def _unwrap(self, item): + return item[0] if self.key is None else item[2] + + def __init__(self, iterable=None, key=None, maxheap=False): + self.lock = Lock() + self.key = key + self.maxheap = maxheap + self.count = count().__next__ + if iterable is None: + self.data = [] + else: + self.data = list(map(self._wrap, iterable)) + heaper = _heapify_max if self.maxheap else heapify + with self.lock: + heaper(self.data) + + def __len__(self): + return len(self.data) + + def push(self, value): + 'Add a value to the heap' + item = self._wrap(value) + pusher = _heappush_max if self.maxheap else heappush + with self.lock: + pusher(self.data, item) + + def pop(self): + 'Remove and return the min (or max) value from the heap' + popper = _heappop_max if self.maxheap else heappop + with self.lock: + result = popper(self.data) + return self._unwrap(result) + + def pushpop(self, value): + 'Same as a push() followed by a pop() but runs a little faster' + item = self._wrap(value) + pushpopper = _heappushpop_max if self.maxheap else heappushpop + with self.lock: + result = pushpopper(self.data, item) + return self._unwrap(result) + + def poppush(self, value): + 'Same as a pop() followed by a push() but runs a little faster' + item = self._wrap(value) + poppusher = _heapreplace_max if self.maxheap else heapreplace + with self.lock: + result = poppusher(self.data, item) + return self._unwrap(result) + +################################################################################ +### Command-line self-test +################################################################################ + if __name__ == "__main__": # Simple sanity test heap = []