Numpy select 索引随机

Question

I have a numpy array that is like this

arr = np.random.randint(2, size=(32, 4, 19))

and I want to output a random index for arr that is equal to 1 in respect to axis=2 . So I want to return an array that is of shape (32, 4, 1) consisting of a random index that is 1. For example, say the first few rows of the array looks like this

array([[[0, 1, 0, ..., 0, 1, 1],
        [0, 0, 1, ..., 1, 1, 0],
        [0, 0, 0, ..., 1, 0, 0],
        [0, 0, 1, ..., 0, 0, 0]],
        ...
       [[0, 1, 1, ..., 1, 1, 1],
        [1, 0, 1, ..., 0, 1, 0],
        [1, 1, 0, ..., 1, 0, 0],
        [0, 0, 0, ..., 0, 1, 1]]])

I want to get something like

array([[[1],[17],[16],[5]], 
       [[3], ...
        ....
       [[7],[4],[7],[11]]]) 

since arr[0,0,1] == 1 and arr[0,1,17] == 1 etc. Can someone please help me

Answer 1

Assuming you have at least one 1 per vector on axis 2, you can multiply each item by a random value and get the argmax , then reshape to add an extra dimension:

np.argmax(arr*np.random.random(size=arr.shape), axis=2)[...,None]

Example output:

[[[ 6], [ 7], [ 9], [ 5]],
 [[10], [15], [13], [16]],
 [[ 1], [ 9], [ 3], [16]],
 [[ 0], [18], [14], [ 0]],
 [[10], [14], [ 1], [11]],
...
]

NB. If the last dimension (19) is extremely large, you might have a bias towards lower indices in case twice the same random value is generated ( argmax will pick the first max), but this is extremely unlikely and will still give a correct result.