在numpy数组中无需使用for循环进行迭代

Question

我需要在numpy数组上进行逻辑迭代，其值取决于其他数组的元素。我编写了下面的代码以澄清我的问题。有没有不使用for循环解决此问题的建议？

Code
a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = np.zeros_like(b)
c[0] = 150
for i in range(1, c.size):
    if a[i] == "b":
        c[i] = c[i-1]
    else:
        c[i] = b[i]

- Elgin Cahangirov

Answer 1

以下是一种方案，结合使用 np.maximum.accumulate 和 np.where 来创建在特定时间间隔内需要停止的阶梯状索引，然后直接从 b 索引即可得到所需的输出。

因此，实现方式如下 -

mask = a!="b"
idx = np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))
out = b[idx]

样例逐步运行 -

In [656]: # Inputs 
     ...: a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
     ...: b = np.array([150, 154, 147, 126, 148, 125])
     ...: 

In [657]: mask = a!="b"

In [658]: mask
Out[658]: array([ True, False,  True,  True, False,  True], dtype=bool)

# Crux of the implmentation happens here :
In [696]: np.where(mask,np.arange(mask.size),0)
Out[696]: array([0, 0, 2, 3, 0, 5])

In [697]: np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))
Out[697]: array([0, 0, 2, 3, 3, 5])# Stepped indices "intervaled" at masked places

In [698]: idx = np.maximum.accumulate(np.where(mask,np.arange(mask.size),0))

In [699]: b[idx]
Out[699]: array([150, 150, 147, 126, 126, 125])

- Divakar

Answer 2

您可以使用更矢量化的方法，如下所示：

np.where(a == "b", np.roll(c, 1), b)

np.where函数根据条件返回符合条件的数组元素。如果条件为True，则从np.roll(c, 1)中获取元素；如果条件为False，则从b中获取元素。np.roll(c, 1)会将c数组中的所有元素向前“滚动”1个位置，这样每个元素都会指向c[i-1]。

这种类型的操作使得numpy变得非常有价值。应尽可能避免使用循环。

- pbreach

---

这是一个简单且简洁的解决方案。但是，为什么它返回[150 0 147 126 0 125]？它没有从“b”数组中取值，其中a[i] =“ b”。 - Elgin Cahangirov

我误读了问题，这只适用于您已经拥有c的所有元素，但是您正在通过循环填充它，因此比这更复杂一些。 - pbreach

Answer 3

如果您不需要环绕边距，那么有一个非常简单的解决方案：

a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = b.copy()  #removes necessity of else case
c[a[:-1]=='b'] = c[a[1:]=='b']

或同样：

a = np.array(['a', 'b', 'a', 'a', 'b', 'a'])
b = np.array([150, 154, 147, 126, 148, 125])
c = b.copy()  #removes necessity of else case
mask = a == 'b'
c[mask[:-1]] = c[mask[1:]]

如果你想要环绕边缘（a[0]=='b'），那么就会变得有点复杂，你需要使用roll或者先用if语句来捕获这种情况。

- DerWeh