如何有条件地更改数组值

Question

我需要"有条件地"向现有的numpy数组添加2个新列,并希望在没有太多复杂性的情况下解决它.

假设我有以下数组:

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

对于这个数组,我需要添加2个新列,这将使数组看起来如下所示:

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

这是它背后的逻辑:

原始数组有2列.输出有4列.(3.和4.列是新的.)

3.列本质上是第二列的增量和,但只有在数组的第一列中相应的值为1时,才会继续添加第2列中的值(来自[:,1])(a: ,0]).例如:

• a [0,2]是2,因为[0,1]是2,我们接受它因为[0,0] = 1
• a [1,2]仍然为"2"因为[1,0] =( - 1)所以我们跳过a [1,1]的值
• a [2,2]变为(2 + 6 =)8.它是[0,2]和[2,2]之和.只要第一列中相应的行值为NOT(-1),我们只对2.列中的值求和.
• a [3,2]保持为8,因为[3,0] =( - 1)因此a [3,1]不加到总和中.

创建4.列是相同的,但这次只需添加第二列(a [:,1])中的值,只要第一列的行值为-1即可.

是否有任何库函数以一种很好的方式支持这样的操作？

Answer 1

你可以使用np.cumsum第一列口罩的产品与第二列,以获得累积和再numpy.stack加入数组:

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)

产量

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

请注意,您必须重塑两者third并fourth匹配尺寸a.