indices.py 文件源码

def combinations(seq, k):
    """ Return j length subsequences of elements from the input iterable.

    This version uses Numpy/Scipy and should be preferred over itertools. It avoids
    the creation of all intermediate Python objects.

    Examples
    --------

    >>> import numpy as np
    >>> from itertools import combinations as iter_comb
    >>> x = np.arange(3)
    >>> c1 = combinations(x, 2)
    >>> print(c1) # doctest: +NORMALIZE_WHITESPACE
       [[0 1]
       [0 2]
       [1 2]]
    >>> c2 = np.array(tuple(iter_comb(x, 2)))
    >>> print(c2) # doctest: +NORMALIZE_WHITESPACE
       [[0 1]
       [0 2]
       [1 2]]
    """
    from itertools import combinations as _combinations, chain
    from scipy.misc import comb

    count = comb(len(seq), k, exact=True)
    res = np.fromiter(chain.from_iterable(_combinations(seq, k)),
                      int, count=count*k)
    return res.reshape(-1, k)