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