python - 如何有条件地更改数组值

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])( [:, 0]).例如:

• a[0, 2] 是 2 因为 a[0, 1] 是 2 我们取它是因为 a[0,0] = 1
• a[1, 2] 保持为 "2"因为 a[1,0]=(-1) 所以我们跳过 a[1, 1] 的值
• a[2, 2] 变成 (2 + 6 =) 8。它是 a[0, 2] 和 a[2, 2] 的和。我们只对第 2 列的值求和，只要第一列中对应的行值不是 (-1)。
• a[3, 2] 保持为 8 因为 a[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]]

请注意，您必须 reshape third 和 fourth 以匹配 a 的尺寸。