from heapq import nlargest as _nlargest, nsmallest as _nsmallest from __builtin__ import min as _min, max as _max from itertools import tee, izip, imap, count from operator import itemgetter __all__ = ['nsmallest', 'nlargest', 'min', 'max'] def nsmallest(n, iterable, key=None): """Find the n smallest elements in a dataset. Equivalent to: sorted(iterable, key=key)[:n] """ if key is None: return _nsmallest(n, iterable) in1, in2 = tee(iterable) it = izip(imap(key, in1), count(), in2) # decorate result = _nsmallest(n, it) return map(itemgetter(2), result) # undecorate def nlargest(n, iterable, key=None): """Find the n largest elements in a dataset. Equivalent to: sorted(iterable, key=key, reverse=True)[:n] """ if key is None: return _nlargest(n, iterable) in1, in2 = tee(iterable) it = izip(imap(key, in1), count(), in2) # decorate result = _nlargest(n, it) return map(itemgetter(2), result) # undecorate def min(iterable, key=None): if key is None: return _min(iterable) it = iter(iterable) try: min_elem = it.next() except StopIteration: raise ValueError('min() arg is an empty sequence') min_k = key(min_elem) for elem in it: k = key(elem) if k < min_k: min_elem = elem min_k = k return min_elem def max(iterable, key=None): if key is None: return _max(iterable) it = iter(iterable) try: max_elem = it.next() except StopIteration: raise ValueError('max() arg is an empty sequence') max_k = key(max_elem) for elem in it: k = key(elem) if k > max_k: max_elem = elem max_k = k return max_elem