Select random non zero elements from each row of a 2d numpy array

Question

I have a 2d array

a = array([[5, 0, 1, 0],
           [0, 1, 3, 5],
           [2, 3, 0, 0],
           [4, 0, 2, 4],
           [3, 2, 0, 3]])

and a 1d array

b = array([1, 2, 1, 2, 2])

which ( b ) tells how many non-zero elements we want to choose from each row of the array a .

For example, b[0] = 1 tells us that we have to choose 1 non-zero element from a[0] , b[1] = 2 tells us that we have to choose 2 non-zero elements from a[1] , and so on.

For a 1d array, it can be done using np.random.choice , but I can't find how to do it for a 2d array, so I have to use a for loop which slows the computation.

I want the result as 2d array as

array([[5, 0, 0, 0],
       [0, 1, 0, 5],
       [2, 0, 0, 0],
       [0, 0, 2, 4],
       [3, 2, 0, 0]])

Here, we have 1 element in row 1, 2 elements in row 2, 1 element in row 3 and so on as given in array b.

Answer 1

It looks like a Competitive Programming problem.

I don't think that you can achieve the results using numpy.random.choice (I may be wrong).

Anyways, think of it like this. To select x number of non-zero elements from a 1D array of size n , it will be of O(n) complexity in the worst case. And, for a 2D array it will be O(n^2) if you follow the same naive approach.

this post is almost similar to your question, but numpy.nonzero also is an O(n^2) function.