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列中的值（来自a [:, 1]）。 a [:, 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 [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]]

请注意，您必须同时重塑third和fourth的形状以匹配a的尺寸。

Answer 2

这是三班轮：

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]])

一种变体要快一些，主要是因为它消除了方便但缓慢的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]])

如何有条件地更改数组值

2 个答案: