Как получить все возможные комбинации тасования, используя numpy

Question

У меня есть массив NumPy длиной N, содержащий X единиц и N-X нулей. Я хочу сгенерировать все возможные комбинации случайных чисел для моего массива.

Например, когда N равен 5, а X равен 2, мне нужны следующие комбинации:

[0, 0, 1, 1, 0]
[1, 1, 0, 0, 0]
[0, 0, 0, 1, 1]
[0, 0, 1, 0, 1]
[1, 0, 0, 1, 0]
[0, 1, 1, 0, 0]
[0, 1, 0, 1, 0]
[1, 0, 1, 0, 0]
[0, 1, 0, 0, 1]
[1, 0, 0, 0, 1]

Is Есть ли прямой способ сделать это в numpy?

Answer 1

Комбинации реализованы в стандартной библиотеке Python, в itertools. Он выбирает X элементов из списка N, поэтому вы можете позволить ему выбирать индексы и устанавливать для них 1:

import numpy as np
from itertools import combinations

N = 5
X = 2

for combination in combinations(range(N), X):
    arr = np.zeros(N)
    arr[list(combination)] = 1
    print(arr)

Вывод:

[1. 1. 0. 0. 0.]
[1. 0. 1. 0. 0.]
[1. 0. 0. 1. 0.]
[1. 0. 0. 0. 1.]
[0. 1. 1. 0. 0.]
[0. 1. 0. 1. 0.]
[0. 1. 0. 0. 1.]
[0. 0. 1. 1. 0.]
[0. 0. 1. 0. 1.]
[0. 0. 0. 1. 1.]

Answer 2

Вот векторизация с itertools.combinations -

from itertools import combinations

def combs(N, X):
    idx = np.array(list(combinations(range(N), X)))
    n = len(idx)
    out = np.zeros((n,N), dtype=int)
    out[np.arange(n)[:,None],idx] = 1
    return out

Образцы прогонов -

In [57]: combs(N=5, X=2)
Out[57]: 
array([[1, 1, 0, 0, 0],
       [1, 0, 1, 0, 0],
       [1, 0, 0, 1, 0],
       [1, 0, 0, 0, 1],
       [0, 1, 1, 0, 0],
       [0, 1, 0, 1, 0],
       [0, 1, 0, 0, 1],
       [0, 0, 1, 1, 0],
       [0, 0, 1, 0, 1],
       [0, 0, 0, 1, 1]])

In [58]: combs(N=5, X=3)
Out[58]: 
array([[1, 1, 1, 0, 0],
       [1, 1, 0, 1, 0],
       [1, 1, 0, 0, 1],
       [1, 0, 1, 1, 0],
       [1, 0, 1, 0, 1],
       [1, 0, 0, 1, 1],
       [0, 1, 1, 1, 0],
       [0, 1, 1, 0, 1],
       [0, 1, 0, 1, 1],
       [0, 0, 1, 1, 1]])

In [59]: combs(N=5, X=4)
Out[59]: 
array([[1, 1, 1, 1, 0],
       [1, 1, 1, 0, 1],
       [1, 1, 0, 1, 1],
       [1, 0, 1, 1, 1],
       [0, 1, 1, 1, 1]])