在 3d numpy 数组中查找特定行的索引

Question

我正在寻找一种在 numpy 中查找 3d 数组中外部特定行的索引的方法。一个示例是在 RBG 图像中查找给定颜色集的所有出现，并获取像素坐标。

这个question 表明in 运算符对数组的行为可能很奇怪，而这个one 更接近但适用于二维数组。

假设我们有一个维度为(x,y,z) 的3d 数组Z，以及我们想要匹配的第三维度行[s0, s1]。

Z = np.zeros((10,20,3), dtype=int)
s0 = np.array([1,2,3])
s1 = np.array([4,5,6])
Z[1,2] = s0
Z[4,5] = s1

我想要所有(x,y)，其中z 等于s0 或s1。

到目前为止， argwhere 返回 s0 中的一个元素在 Z 中的每个匹配项：

> np.argwhere(s0 == Z) 
array([[1, 2, 0],
       [1, 2, 1],
       [1, 2, 2]])

in1d 返回一个布尔型一维数组，其中 s0 或 s1 中的元素匹配：

> np.in1d(Z, [s0,s1])

如果我尝试解开的方式：

> Zravel = np.ascontiguousarray(a).view([('', a.dtype)] * a.shape[-1]).ravel()
> np.all(np.in1d(Zravel, [s0, s1]) == False)

所有元素都是False。

有什么想法吗？

Answer 1

np.in1d 会展平其输入。所以，你可以给它Z和s0、s1的堆叠版本，给我们一个布尔数组，它可以被重新塑造成一个与Z形状相同的数组。然后，您需要检查其中的所有 TRUE 行是否有匹配的索引。实现看起来像这样 -

S = np.row_stack((s0,s1))
out = np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))

你也可以用broadcasting这样解决-

out = np.where(((Z == S[:,None,None,:]).all(3)).any(0))

如果您希望将输出堆叠在一个数组中 -

outarr = np.column_stack((out))

要创建S，可以将np.row_stack 替换为np.concatenate，这样可能会更快，像这样-

S = np.concatenate((s0,s1)).reshape(-1,s0.size)

示例运行 -

In [145]: Z = np.zeros((10,20,3), dtype=int)
     ...: s0 = np.array([1,2,3])
     ...: s1 = np.array([4,5,6])
     ...: Z[1,2] = s0
     ...: Z[4,5] = s1
     ...: 

In [146]: np.where(((Z == S[:,None,None,:]).all(3)).any(0))
Out[146]: (array([1, 4]), array([2, 5]))

In [147]: np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))
Out[147]: (array([1, 4]), array([2, 5]))

In [148]: np.column_stack((np.where(((Z == S[:,None,None,:]).all(3)).any(0))))
Out[148]: 
array([[1, 2],
       [4, 5]])

In [149]: np.column_stack((np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))))
Out[149]: 
array([[1, 2],
       [4, 5]])