Search code examples
pythonnumpyrandomprobability

random.choice for 2D array, bigger numbers with higher probability?

Is there a way to randomly pick n-items from every row in a 2D array with the higher probability picking the bigger values w/o using a LOOP

random.choice() works only for 1D array...

F.e. if i have :

 q = np.random.random((10,10))

i can pick the max-2 in every row like this :

  np.sort(q,axis=1)[:,-2:]

what I want instead is to pick randomly 2 not always the max, but with higher probability the bigger numbers..

here is how you get single row with probabilities :

np.random.choice(q[0,:], p=q[0,:]/q[0,:].sum())
Solution

  • You can use apply_along_axis:

    q = np.random.random((10,10))

def choice(row, n, replace=False):
    return np.random.choice(row, size=n, p=row/row.sum(), replace=replace)

np.apply_along_axis(func1d=choice, axis=1, arr=q, n=2)

    I don't know what array do you have, but you should probably check that row.sum() is not 0 to avoid errors in computation of p=row/row.sum().