获取 ndarray 中 N 个最高值的索引 [英] Get the indices of N highest values in an ndarray

问题描述

考虑形状为 100x100x100 的直方图，我想找到 2 个最高值 a 和 b，以及它们的索引 (a1, a2, a3) 和 (b1, b2, b3)，例如:

hist[a1][a2][a3] = ahist[b1][b2][b3] = b

我们可以使用 hist.max() 轻松获得最高值，但是我们如何获得 ndarray 中的 X 个最高值?

我知道人们通常使用 np.argmax 来检索值索引，但在这种情况下:

hist.argmax().shape = () # 单值对于范围内的 i (3):hist.argmax(i).shape = (100, 100)

如何获得形状 (3)，即每个维度一个值的元组?

解决方案

您可以使用 numpy.argpartition 首先在数组的扁平版本上获取顶部 k 项的索引，然后您可以根据数组的形状使用转换这些一维索引numpy.unravel_index:

<预><代码>>>>arr = np.arange(100*100*100).reshape(100, 100, 100)>>>np.random.shuffle(arr)>>>指数 = np.argpartition(arr.flatten(), -2)[-2:]>>>np.vstack(np.unravel_index(indices, arr.shape)).T数组([[97, 99, 98],[97, 99, 99]]))>>>arr[97][99][98]999998>>>arr[97][99][99]999999

Considering an histogram of shape 100x100x100, I would like to find the 2 highest values a and b, and their indices (a1, a2, a3) and (b1, b2, b3), such as:

hist[a1][a2][a3] = a
hist[b1][b2][b3] = b

We can easily get the highest value with hist.max(), but how can we get the X highest values in a ndarray?

I understand that one normally uses np.argmax to retrieve the value indices, but in that case:

hist.argmax().shape = ()  # single value
for i in range(3):
    hist.argmax(i).shape = (100, 100)

How can I get a shape (3), a tuple with one value per dimension?

解决方案

You can use numpy.argpartition on flattened version of array first to get the indices of top k items, and then you can convert those 1D indices as per the array's shape using numpy.unravel_index:

>>> arr = np.arange(100*100*100).reshape(100, 100, 100)
>>> np.random.shuffle(arr)
>>> indices =  np.argpartition(arr.flatten(), -2)[-2:]
>>> np.vstack(np.unravel_index(indices, arr.shape)).T
array([[97, 99, 98],
       [97, 99, 99]])
)
>>> arr[97][99][98]
999998
>>> arr[97][99][99]
999999

这篇关于获取 ndarray 中 N 个最高值的索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！