python - 获取任意形状的多维 numpy 数组的索引并将其用于另一个数组

Question

我有两个多维数组a并且b形状相同：

print(a.shape) -> (100, 20, 3, 3)
print(b.shape) -> (100, 20, 3, 3)

现在我只想从a符合特定条件的情况中获取值，例如：

wanted_value = 5
result_a = a[a[:, :, :, 0] == wanted_value]

但是，这很好用，现在假设我想从另一个索引数组b中获取数据，这些索引对应于 array 中找到的元素的索引a。是否有可能以某种方式使用我可以从搜索数组中获得的索引在数组a上的某个条件下从数组b中获取b与数组具有的相同形状的相同索引相对应的数据a？我需要这个，因为我无法在数组中搜索b我需要的条件，但数组的数据/元素分别b对应于数组a。

这可能吗？我试过这样的事情：

indices = np.where(a[a[:, :, :, 0] == wanted_value])
results_b = b[indices]

标签： pythonpython-3.xnumpynumpy-ndarray

Answer 1

是的。它比你想象的还要简单：

b[a[:, :, :, 0] == wanted_value]

它的工作方式是它a[:, :, :, 0] == wanted_value自己创建一个带有 的 Numpy 数组，dtype=bool在文档中称为掩码。您可以将其应用于任何相同形状的数组。