jeudi 24 juin 2021

Is there a short equivalent of matlab's nchoosek function in numpy or scipy?

In matlab nchoosek gives you all possible combinations k values taken from a set of size n:

 nchoosek(1:4,2)
ans =

   1   2
   1   3
   1   4
   2   3
   2   4
   3   4

The closest solution I came up with in Python was:

np.vstack([rng.choice(np.arange(5),size=3, replace=False, shuffle=False) for i in range(int(scipy.special.comb(5,2)))])

Which gives:

array([[0, 1, 4],
       [1, 3, 4],
       [0, 3, 4],
       [1, 2, 4],
       [0, 2, 4],
       [2, 1, 4],
       [0, 3, 4],
       [1, 2, 4],
       [2, 3, 1],
       [0, 1, 4]])

Isn't there a shorter and probably more efficient way to do this in python?




Aucun commentaire:

Enregistrer un commentaire