python - 对从 np.where 获得的索引应用偏移量

Question

我有一个 3d numpy 数组，我获取满足特定条件的索引，例如:

a = np.tile([[1,2],[3,4]],(2,2,2))
indices = np.where(a == 2)

对于这个索引，我需要应用一个偏移量，例如 (0, 0, 1)，并查看是否满足另一个条件。

类似这样的事情:

offset = [0, 0, 1]
indices_shift = indices + offset

count = 0
for i in indices_shift:
    if a[i] == 3:
       count += 1

在此示例中，偏移量为 (0,0,1)，索引如下所示:

indices = (array([0, 0, 0, 0, 1, 1, 1, 1], dtype=int64), array([0, 0, 2, 2, 0, 0, 2, 2], dtype=int64), array([1, 3, 1, 3, 1, 3, 1, 3], dtype=int64))

我认为添加偏移量结果应该是这样的:

indices_shift = (array([0, 0, 0, 0, 1, 1, 1, 1], dtype=int64), array([0, 0, 2, 2, 0, 0, 2, 2], dtype=int64), array([2, 4, 2, 4, 2, 4, 2, 4], dtype=int64))

有什么简单的方法可以做到这一点吗？

谢谢。

Answer 1

这是一种方法 -

idx = np.argwhere(a == 2)+[0,0,1]
valid_mask = (idx< a.shape).all(1)
valid_idx = idx[valid_mask]
count = np.count_nonzero(a[tuple(valid_idx.T)] == 3)

步骤:

• 获取与 2 匹配的索引。使用np.argwhere在这里将它们放入一个漂亮的二维数组中，每列代表一个轴。另一个好处是，这使得处理具有通用维数的数组变得通用。然后，以广播的方式添加offset。这是idx。

• 在idx中的索引中，很少有超出数组形状的无效索引。因此，获取有效掩码 valid_mask 以及其中的有效索引 valid_idx。

• 最后用这些索引到输入数组中，与 3 进行比较并计算匹配的数量。