Generate all n choose k binary vectors python

Viewed 156

Is there any efficient way (numpy style) to generate all n choose k binary vectors (with k ones)? for example, if n=3 and k=2, then I want to generate (1,1,0), (1,0,1), (0,1,1).

Thanks

1 Answers

I do not know how efficient this is, but here is a way:

from itertools import combinations
import numpy as np

n, k = 5, 3


np.array(
    [
        [1 if i in comb else 0 for i in range(n)]
        for comb in combinations(np.arange(n), k)
    ]
)
>>>
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.]])
Related