Python 在多维数组上使用numpy where，但需要控制索引

如果3D数组的元素超过某个阈值，我需要修改它们。修改基于另一个数组的相关元素。更具体地说：

A_ijk = A_ijk if A_ijk < threshold value

= (B_(i-1)jk + B_ijk) / 2, otherwise

where提供了我需要的大部分功能，但我不知道如何在没有显式循环的情况下迭代第一个索引。下面的代码执行我想要的操作，但使用循环。有更好的办法吗？假设A和B是相同的形状

for i in xrange(A.shape[0]):
    A[i] = numpy.where(A[i] <= threshold, A[i], (B[i - 1] + B[i]) / 2)

---

为了解决下面的问题：A的前几行保证低于阈值。这可以防止i索引循环到A的最后一个条目。

您可以使用布尔索引替换A中高于阈值的元素，从而将操作矢量化。必须小心谨慎，因为对应于B[i-1]+B[i]/2的辅助数组在第一维度上的大小比A小一个，因此我们必须明确忽略A的第一行，因为它们都低于阈值，如问题中所述：

import numpy as np

# some dummy data
A = np.random.rand(3,4,5)
B = np.random.rand(3,4,5)
threshold = 0.5
A[0,:] *= threshold # put the first dummy row below threshhold
mask = A[1:] > threshold # to be overwritten, shape (2,4,5)

replace = (B[:-1] + B[1:])/2 # to overwrite elements in A from, shape (2,4,5)

# replace corresponding elements where `mask` is True
A[1:][mask] = replace[mask]

---

您可以使用where直接索引到ndarray：

a   = np.random.rand(4,3,10)
b   = np.zeros(a.shape)
idx = np.where(a < .1)
print(a)
a[idx] = b[idx]
print(a)

如果需要for循环，只需将ravel转换为索引并更新即可

a   = np.random.rand(4,3,10)
b   = np.zeros(a.shape)
idx = [np.ravel_multi_index(i, a.shape) for i in zip(*np.where(a < .1))]
print(a, idx)
for i in idx:
    a.ravel()[i] = b.ravel()[i]
print(a)

---

这个包起来了吗？例如，如果i==0，那么B[i-1]应该引用数组的最后一行？如果A[i，…]同时引用B[i，…]和B[i-1，…]，那么A和B不能具有相同的形状。我的意思是它可以，但是你确定这个循环代码做了你想要它做的事情吗？谢谢！这会奏效的。如上所述，已知A的第一行低于阈值，因此来自B矩阵的掩码值从不循环around@JimParker谢谢，我错过了那个细节，那时我已经开始写答案了。我更新了答案，在这个附加假设下，它稍微简单一些，不需要辅助数组。