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

Question

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

这个question表明in运算符在数组中的行为会很奇怪，这个one更接近于二维数组。在

假设我们有一个3d数组Z，其维数为(x,y,z)，而{}是我们想要匹配的第三维行。在

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或{}。在

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

^{pr2}$

in1d返回布尔1D数组，其中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

^{}将展平其输入。因此，您可以给它提供Z和{}，s1的堆叠版本，给我们一个布尔数组，它可以被重塑成与Z相同形状的数组。然后，您需要检查其中是否有匹配索引的所有真行。执行应该是这样的-

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

您也可以使用^{}来解决它，如下-

^{pr2}$

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

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]])