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

Posted 2023-02-23

【中文标题】查找对应于另一个数组中沿轴的最小值的数组

【英文标题】：Find array corresponding to minimal values along an axis in another array

【发布时间】：2017-05-10 20:52:04

【问题描述】：

我有两个三维数组，a 和 b，并且想要找到 b 的二维子数组，其中元素 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 数组，我必须进行大量的整形操作——我想知道是否有更简单的方法。

【参考方案1】：

这是一种 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])

【讨论】：

将这些重新整形为 (5,2)，结果应该类似于另一个答案建议的ogrid 的mgrid 版本。

@hpaulj 是的，不幸的是我不熟悉那个功能。

【参考方案2】：

您可以使用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 列表。

【讨论】：

只是为了检查：n.ogrid[0:x, 0:y] + [c] 将第三维添加到 ogrid 数组中，其中包含 c 的元素，对吗？所以这似乎与花哨的索引答案非常相似（我已经接受了，因为它更短，但这很有趣 - 谢谢！）。

@mzzx 基本上相似的答案，因为两者都创建广播范围数组以索引到数据数组。不同之处仅在于创建这些范围数组的方式。我的帖子明确地进行广播，而这篇文章是在函数的帮助下完成的。

@mzzx 是的，它们是等价的！如果你喜欢简短，你也可以写成一行：arr[np.ogrid[0:arr.shape[0], 0:arr.shape[1]] + [c]] :-)

【参考方案3】：

这是一种使用fancy-indexing的方法-

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