Python, random sample from a vector's non-zero indices

Question

I have a 1-d vector called visited. I want to find randomly sample k indices of its zero position.

This is what I do now:

random.sample(np.where( visited == 0)[0]) , k) 

Is there a better way of doing it?

Thanks

Answer 1

The way you're doing it is sound. However, you could use the more intuitive nonzero function:

random.sample(visited.nonzero(), k)

EDIT:

As to the second question in you comment, you can inverse the "zeroness" of you array: visited==0. You get:

random.sample((visited==0).nonzero(), k)