Fastest way to Select a random number from each row padded numpy array (excluding the pad) and number of non padded values, using numpy operations

Viewed 155

I have a 2D numpy array, each row is padded with (with -1 for the example below).

For each row, I want to pick a random number, excluding the padding, and also get the number of non-padded values for each row, using only numpy operations.

Here is a minimal example. I picked -1 for the pad, but the pad can by any negative int.

import numpy as np
numList = [[0, 32, 84, 93, 1023, -1], [0, 23, 33, 45, -1, -1], [0, 10, 15, 21, 24, 25], [0, 23, -1, -1, -1, -1], [0 , 13, 33, 34, -1, -1]]
numArray = np.array(numList)
numArray

array([[   0,   32,   84,   93, 1023,   -1],
       [   0,   23,   33,   45,   -1,   -1],
       [   0,   10,   15,   21,   24,   25],
       [   0,   23,   -1,   -1,   -1,   -1],
       [   0,   13,   33,   34,   -1,   -1]])

For the lengths, the output should look something like this

LengthsResults
[5, 4, 6, 2, 4]. 

And here's an example output for picking a random non-pad number for each row.

randomNonPad
[84, 45, 0, 0, 34]

Edit:

I was looking at np.where, which lets you filter out parts of your numpy array on a conditional, and numpy random choice, which lets you pick a random number for an array. I'm not sure what to do with np.where though, it seems that you can change it to something, but I'm not sure what yet, or even if it's the right approach. For python, you could start with a list, and append it to any length, but for numpy you need to establish the array length ahead of time.

2 Answers

The index of the negative number in the row, which is also the length of non-padded elements, is most simply gotten by

lengths = np.argmin(numArray, axis=1)

This assumes that the padding number is the same for all elements within the row. This won't work properly for rows with no negative numbers, so you can fix it with:

lengths[np.take_along_axis(numArray, lengths.reshape(-1, 1), axis=1).ravel() >= 0] = numArray.shape[1]

You can now use this information to generate an array of random indices into your rows:

indices = np.random.randint(lengths)

And apply the index to get the corresponding elements:

result = np.take_along_axis(numArray, indices.reshape(-1, 1), axis=1)

While cleaning up the lengths array is likely the faster option, a shorter expression might be something like

lengths = np.where(np.any(numArray < 0, axis=1), np.argmin(numArray, axis=1), numArray.shape[1])

Also, if your padding number is not a consistent negative number, you will need to replace np.argmin(numArray, axis=1) with either np.argmax(numArray < 0, axis=1), or np.argmin(numArray >= 0, axis=1), regardless of which approach you use to compute lengths.

Note - this probably overlaps with @Mad's answer; I'll leave it in case the alternative explanation clears up some point of confusion.

In [32]: numList = [[0, 32, 84, 93, 1023, -1], [0, 23, 33, 45, -1, -1], [0, 10, 15, 21, 2
    ...: 4, 25], [0, 23, -1, -1, -1, -1], [0 , 13, 33, 34, -1, -1]] 
    ...: numArray = np.array(numList)                                                    
In [33]: numArray                                                                        
Out[33]: 
array([[   0,   32,   84,   93, 1023,   -1],
       [   0,   23,   33,   45,   -1,   -1],
       [   0,   10,   15,   21,   24,   25],
       [   0,   23,   -1,   -1,   -1,   -1],
       [   0,   13,   33,   34,   -1,   -1]])

number of pads per row:

In [34]: np.sum(numArray==-1, axis=1)                                                    
Out[34]: array([1, 2, 0, 4, 2])

number of non-pad per row:

In [35]: np.sum(numArray!=-1, axis=1)                                                    
Out[35]: array([5, 4, 6, 2, 4])

I don't know if assuming the pad values are all at the end makes this any more efficient or not. The sample's a bit small to make good timings.

picking a random non-pad from each row, the obvious first attempt is a row list comprehension:

In [40]: [np.random.choice(row[row!=-1]) for row in numArray]                            
Out[40]: [32, 0, 0, 23, 34]

Alternatively working from the lengths (above) (and assuming tail padding) we could pick a random index for each row:

In [46]: [np.random.choice(i) for i in Out[35]]                                          
Out[46]: [1, 2, 1, 0, 1]
In [47]: numArray[np.arange(numArray.shape[0]), [np.random.choice(i) for i in Out[35]]]  
Out[47]: array([93, 45, 21, 23, 13])

In hat tip to @Mad, randint accepts a list/array of range values, the choice comprehension can be replaced with:

In [49]: np.random.randint(Out[35])                                                      
Out[49]: array([3, 1, 2, 1, 1])
In [50]: numArray[np.arange(numArray.shape[0]), np.random.randint(Out[35])]              
Out[50]: array([ 0, 23, 24,  0,  0])
Related