N-D version of itertools.combinations in numpy

by

You can use itertools.combinations() to create the index array, and then use NumPy’s fancy indexing:

import numpy as np
from itertools import combinations, chain
from scipy.special import comb

def comb_index(n, k):
    count = comb(n, k, exact=True)
    index = np.fromiter(chain.from_iterable(combinations(range(n), k)), 
                        int, count=count*k)
    return index.reshape(-1, k)

data = np.array([[1,2,3,4,5],[10,11,12,13,14]])

idx = comb_index(5, 3)
print(data[:, idx])

output:

[[[ 1  2  3]
  [ 1  2  4]
  [ 1  2  5]
  [ 1  3  4]
  [ 1  3  5]
  [ 1  4  5]
  [ 2  3  4]
  [ 2  3  5]
  [ 2  4  5]
  [ 3  4  5]]

 [[10 11 12]
  [10 11 13]
  [10 11 14]
  [10 12 13]
  [10 12 14]
  [10 13 14]
  [11 12 13]
  [11 12 14]
  [11 13 14]
  [12 13 14]]]

More Related Contents:

  1. Permutations between two lists of unequal length
  2. itertools product speed up
  3. Convert a 3D array to RGB valued array [closed]
  4. Unable to allocate array with shape and data type
  5. Replacing Numpy elements if condition is met
  6. How to make a multidimension numpy array with a varying row size?
  7. What does the term “broadcasting” mean in Pandas documentation?
  8. Save / load scipy sparse csr_matrix in portable data format
  9. figure of imshow() is too small
  10. Binary random array with a specific proportion of ones?
  11. Difference between np.random.seed() and np.random.RandomState()
  12. Simple wrapping of C code with cython
  13. Loading UTF-8 file in Python 3 using numpy.genfromtxt
  14. What does a colon and comma stand in a python list?
  15. convert string representation of array to numpy array in python
  16. How to share numpy random state of a parent process with child processes?
  17. numpy reverse multidimensional array
  18. Fastest pairwise distance metric in python
  19. How do I get Pylint to recognize NumPy members?
  20. python tilde unary operator as negation numpy bool array
  21. Form a big 2d array from multiple smaller 2d arrays
  22. Sieve of Eratosthenes – Primes between X and N
  23. How to shade region under the curve in matplotlib
  24. Randomly select from numpy array
  25. How to use NumPy array with ctypes?
  26. Can’t import cv2; “DLL load failed”
  27. A numpy array unexpectedly changes when changing another one despite being separate
  28. Computing cross-correlation function?
  29. Annotate Time Series plot
  30. Numpy `ValueError: operands could not be broadcast together with shape …`

Leave a Comment