Question

我有两个三维数组a和b，并希望找到b的2D子阵列，其中a的元素沿着第三轴具有最小值，即

a=n.random.rand(20).reshape((5,2,2))
b=n.arange(20).reshape((5,2,2))
c=n.argmin(a,2) #indices with minimal value of a
d=n.zeros_like(c) #the array I want
for i in range(5):
  for j in range(2):
     d[i,j] = b[i,j,c[i,j]]

有没有办法可以在没有双循环的情况下获得这些值？

我知道这个答案： replace min value to another in numpy array 但如果我希望这对我的3D阵列起作用，我就必须进行大量的整形操作 - 而且我想知道是否有更简单的东西。

Answer 1

您可以使用np.ogrid为其他维度创建网格：

x, y, z = arr.shape  # assuming your array is named "arr"
xyz = np.ogrid[0:x, 0:y] + [c]  # c is your second axis index (the argmin)
arr[xyz]

如果它不是最后一个轴，那么你可以简单地使用insert因为ogrid返回一个包含索引的普通python列表。

Answer 2

这是一种Numpythonic方式：

In [83]: x, y, z = a.shape
In [84]: b[np.repeat(np.arange(x), y), np.tile(np.arange(y), x), c.ravel()].reshape(x, y)

此处np.repeat(np.arange(x), y)将为您提供第一轴的相应索引。

In [86]: np.repeat(np.arange(x), y)
Out[86]: array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4])

np.tile(np.arange(y), x)将为您提供第二轴的相应索引。

In [87]: np.tile(np.arange(y), x)
Out[87]: array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])

对于第三个，你可以使用扁平的c形状。

In [88]: c.ravel()
Out[88]: array([1, 0, 1, 1, 1, 0, 1, 1, 0, 0])

Answer 3

以下是使用fancy-indexing -

的方法

m,n,r = b.shape
d_out = b[np.arange(m)[:,None],np.arange(n),c]

查找沿另一个数组中的轴的最小值对应的数组

3 个答案: