Search code examples
pythonnumpymasked-array

Writing an efficient code for applying a reverse operation using masks in NumPy


Firstly, I have a compressed numpy array c and a mask m that was used to generate c from a full array a.

I want to output a reconstructed array b of the same shape as the original array a, but with the results from the compressed array. The following code works for this, but I don't know how to make it efficient. Any guidance will be appreciated

import numpy as np
a = np.asarray((1, 2, 3, 4, 5, 6, 7, 8, 9))
m = np.array((True,True,True,True,False,False,True,True,True))
c = np.ma.compressed(np.ma.masked_where(m==False, a))

i=0
j=0
b = np.zeros(a.size)
while (i<a.size):
    if (m[i] == True):
        b[i] = c[j]
        j = j+1
    i = i+1
b

which results in:

array([1., 2., 3., 4., 0., 0., 7., 8., 9.])

Solution

  • You can use boolean indexing:

    b = np.zeros_like(m, dtype=float) # change dtype to `int` if that's what you want.
    
    b[m] = c
    

    Output

    array([1., 2., 3., 4., 0., 0., 7., 8., 9.])