当多个选择时在numpy数组中查找随机出现

Question

I wrote this code to achieve the goal of returning a random value from a list of elements that are matching a predicate condition:

N=<int>
sampl = np.random.randint(low=0, high=N+1, size=(10,))
xs = np.where(sampl == 1)
ys = np.array([tuple(x) for x in xs], dtype=int)[0]
x = np.random.choice(ys)

Ex: If I run the code with N=2 and I am looking for only 1 s in the array:

    sampl = np.random.randint(low=0, high=N+1, size=(10,))

--> sampl = [2 1 0 0 0 1 0 0 2 1]

    xs = np.where(sampl == 1)

--> [2 1 0 0 0 1 0 0 2 1]  # Positions 1, 5, 9 are of interest. 
       ^       ^       ^ 

    ys = np.array([tuple(x) for x in xs], dtype=int)[0]

--> ys = [1 5 9] # Put them in an array. 

    x = np.random.choice(ys)

--> x = 9 # Pick a random one and return it

It works but it's not concise and I ran into a few issues trying to make it more elegant.

• numpy.where() returns a tuple when passing nothing else but a condition. I tried passing x=sampl but the runtime complains saying that the function doesn't take params (it does when I inspect the code).
• Again, making a numpy array from a tuple forces me to return the first element. That's error-prone when testing for edge-cases (such as no value found by predicate.)

Do you have any suggestions to improve this code? I want to stick to numpy/pandas as the arrays will become very big.

Answer 1

Probably the most elegant way I could think of is to randomly shuffle your array, and then pull off the first occurrence. That should be pretty concise.

So something like:

np.random.shuffle(sampl)
x = np.ravel(np.where(sampl==1))[0]

or, like you suggested, without shuffling, that would look something like

x = np.random.choice(np.ravel(np.where(sampl==1)))

On second thought, I guess the choice method will be infinitely faster than shuffling.

The next issue is the edge cases. How to handle this depends on what you expect the default behavior to be. If you expect that in most cases the condition will turn up at least one hit, then you should handle the case when there is not a hit with an exception:

try: 
   x = np.random.choice(np.ravel(np.where(sampl==1)))
except: 
   # TODO
   pass

I would highly suggest doing this unless you rarely find a hit. But don't take my word for it... time it yourself.

The other option would be to put in a condition that explicitly checks that

np.size( np.where(sampl==1) ) > 0

before continuing. However, I would guess that that approach is slower than the try...except approach.