Updating nonzero values in an array, A, with the values of a different array, B, whose size is equal to number of nonzero values in A

Question

I have an array, A. Its length may vary but it’s always filled with 1’s and/or 0's.

A = np.array([1,0,1,0])

A gets passed into a function that produces array B.

B = np.array([0.75, 0.25])

B’s length is always equal to the number of 1's in A.

How can I most efficiently update A (or create a new array) that equals

[0.75, 0, 0.25, 0]

in this example? My hope is for it to work with any size array A that meets the constraints I’ve laid out for B above. The first value of B is always the first 1 in A and so on.

I got it to work by converting them to lists and looping

pos = 0
for i in range(len(a)):
    if a[i] == 1:
        a[i] = b[pos]
        pos += 1

But I’m hoping for a better solution because I’ll have to do this a lot and A and B could potentially be much larger.

Answer 1

In [1]: A = np.array([1,0,1,0])

First you want to identify the nonzero elements. There are several ways:

In [2]: A.astype(bool)
Out[2]: array([ True, False,  True, False])
In [3]: A>0
Out[3]: array([ True, False,  True, False])
In [4]: np.nonzero(A)
Out[4]: (array([0, 2]),)

Any of these can be used to select those elements; both for getting and setting:

In [5]: A[np.nonzero(A)]
Out[5]: array([1, 1])
In [6]: A[np.nonzero(A)] = [2,3]
In [7]: A
Out[7]: array([2, 0, 3, 0])
In [8]: A[[0,2]]
Out[8]: array([2, 3])

But if we try to assign the two float values, they get truncated:

In [9]: A[[0,2]] = [.25, .75]
In [10]: A
Out[10]: array([0, 0, 0, 0])

A needs to be float dtype:

In [11]: A1=A.astype(float)
In [12]: A1[[0,2]] = [.25, .75]
In [13]: A1
Out[13]: array([0.25, 0.  , 0.75, 0.  ])