Random index of a non zero value in a numpy array

Question

I hope this hasn't been asked before. Not sure how I can reformulate the question. This post is quite similar but still not exactly what I'm looking for.

I have a numpy array with a variable length (something between 4-12). The values of the array are either 0 or 1. I'd like to get the index of a random sample from the array, which is not 0.

My idea was to do something like this:

def sample(self):
        flag = 0
        while flag == 0:
            sampleValue = randint(0, len(myArray())-1)
            flag = myArray()[sampleValue]
        return sampleValue

But that's not really good code and will eventually get trapped in a never ending while loop. Of course I could improve this. Chances are high there's a much more beautiful way of doing this in python :)

Answer 1

You can get the indices where an array a is nonzero with numpy.nonzero

>>> a = np.array([1, 0, 1, 0, 1, 1, 1, 0, 0])
>>> idx_nonzero, = np.nonzero(a)
>>> 
>>> idx_nonzero
array([0, 2, 4, 5, 6])

and then choose one of these values with numpy.random.choice.

>>> np.random.choice(idx_nonzero)
5

Answer 2

You can first generate the list of indices that are ones with:

my_idcs = np.where(my_array)[0]

then you can randomly obtain one of the indices, for example with random.choice:

from random import choice

my_idx = choice(np.where(my_array)[0])

this will thus return a single index where the corresponding value my_array[my_idx] is 1.

Note that if there are no 1s in my_array, then this will raise an IndexError.