CODEWITHSUNDEEP
pythonarraysperformancenumpyindexing
sccrthlt
sccrthlt

Reputation: 4344

Get index of random true value PER COLUMN in numpy 2-d array without for loop

I would like to obtain the index of a random true value per column in a numpy 2-d array using the most efficient means possible (i.e. without a python for loop).

For example, given the following 2-d numpy.ndarray:

np.array(
    [[True, False, False],
     [True, True, True],
     [False, True, False]]
)

Provide a 1-d numpy.ndarray that gives the random row index positions of true values per column (i.e. first axis):

numpy.ndarray([0, 2, 1])

By counter-example, the following would NOT be correct:

numpy.ndarray([2, 2, 1])

because, the 1st column has a non-True value in the third row.

Upvotes: 1

Views: 151

Answers (1)

Dani Mesejo
Dani Mesejo

Reputation: 61910

You could do the following:

import numpy as np

a = np.array(
    [[True, False, False],
     [True, True, True],
     [False, True, False]]
)

result = np.argmax(a * np.random.randint(1, 100, size=a.shape), axis=0)
print(result)

Output

[1 2 1]

The idea is that np.argmax will always choose a position where the array is True, then which one to choose randomly is done by the random.randint function. To each True value it will assign a random integer and you choose the index of the maximum.

Upvotes: 3

Related Questions