Search code examples
pythonpython-3.xnumpy

Generate all n choose k binary vectors python

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

Solution

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

    from itertools import combinations
import numpy as np

def n_choose_k_vectors(n: int, k: int):
    return np.array(
        [
            [int(i in comb) for i in range(n)]
            for comb in combinations(np.arange(n), k)
        ]
    )

n_choose_k_vectors(n=5, k=3)
>>>
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.]])