how to conditionally change array values

Question

I need to add 2 new columns to an existing numpy array "conditionally" and would like to solve it without too much complexity.

Assume that I have the following array:

a = np.array([[1,  2], 
              [-1, 4],
              [1,  6],
              [-1, 8]])

To this array, I need to add 2 new columns which will make the array look like the following:

a = np.array([[1,  2, 2, 0 ], 
              [-1, 4, 2, 4 ],
              [1,  6, 8, 4 ],
              [-1, 8, 8, 12]])

Here is the logic behind it:

The original array has 2 columns. And the output has 4 columns. (3. and 4. columns are new.)

The 3. column is essentially the incremental sum of the second column but you only keep adding the values in column 2 (from a[:, 1]) only if the corresponding value is 1 in the first column of the array (a[:, 0]). For instance:

• a[0, 2] is 2 because a[0, 1] is 2 and we take it because a[0,0] = 1
• a[1, 2] remains as "2" because a[1,0]=(-1) so that we skip the value of a[1, 1]
• a[2, 2] becomes (2 + 6 =) 8. It is the sum of a[0, 2] and a[2, 2]. We only sum the values from the 2. column as long as the corresponding row value is NOT (-1) in the first column.
• a[3, 2] remains as 8 because a[3, 0] = (-1) so that a[3, 1] is not added to the sum.

Creation of the 4. column is the same but this time you just add the values from the second column (a[:, 1]) as long as the row value from the first column is -1.

Is there any library function supporting such an operation in a nice way?

Answer 1

You could use np.cumsum on the product of a mask in the first column with the second column to get the cumulative sum and then numpy.stack to join the arrays:

import numpy as np

a = np.array([[1, 2],
              [-1, 4],
              [1, 6],
              [-1, 8]])

mask = a[:, 0] == 1
third = np.cumsum(a[:, 1] * mask).reshape(-1, 1)

mask = a[:, 0] == -1
fourth = np.cumsum(a[:, 1] * mask).reshape(-1, 1)

result = np.hstack((a, third, fourth))

print(result)

Output

[[ 1  2  2  0]
 [-1  4  2  4]
 [ 1  6  8  4]
 [-1  8  8 12]]

Note that you must reshape both third and fourth to match the dimensions of a.

Answer 2

Here is a three-liner:

a = np.array([[1,  2], 
              [-1, 4],
              [1,  6],
              [-1, 8]])

b = np.zeros(np.arange(2) + a.shape, a.dtype)
np.put_along_axis(b, *a[None].T, 1)
np.c_[a, b[:, 1:].cumsum(0)]
# array([[ 1,  2,  2,  0],
#        [-1,  4,  2,  4],
#        [ 1,  6,  8,  4],
#        [-1,  8,  8, 12]])

And a variant that is a bit faster mostly because it does away with the convenient but slow np.c_

b = np.zeros_like(a)
np.put_along_axis(b, a[:,:1].clip(None, 0), a[:,1:], 1)
np.concatenate([a, b.cumsum(0)], 1)
# array([[ 1,  2,  2,  0],
#        [-1,  4,  2,  4],
#        [ 1,  6,  8,  4],
#        [-1,  8,  8, 12]])