如何从多维数组中获取k个最大值的索引

2024-03-29 13:28:55 发布

您现在位置:Python中文网/ 问答频道 /正文

我在StackOverflow上查了几个问题,但没有找到相应的答案。我想从a numpyndarray获取k个最大值的索引。This link讨论的是相同的,但对于1D数组。^二维数组的{}导致元素按行排序。i、 e

Note: array elements are not unique.

输入:

^{pr2}$

还有

import numpy as np
n = np.arange(9).reshape(3,3)
s = n.argsort(axis=None)
>>>s
array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=int32)

但我在这里失去了数组结构,无法赎回元素的原始索引。在

任何帮助都是感激的。在


Tags: 答案元素排序nplinknot数组elements
1条回答
网友
1楼 · 发布于 2024-03-29 13:28:55

使用^{}和{a2}对ndarrays-

def k_largest_index_argpartition_v1(a, k):
    idx = np.argpartition(-a.ravel(),k)[:k]
    return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argpartition_v2(a, k):
    idx = np.argpartition(a.ravel(),a.size-k)[-k:]
    return np.column_stack(np.unravel_index(idx, a.shape))

def k_largest_index_argsort(a, k):
    idx = np.argsort(a.ravel())[:-k-1:-1]
    return np.column_stack(np.unravel_index(idx, a.shape))

argpartition讨论两个版本

k_largest_index_argpartition_v1和{}之间的区别在于我们如何使用argparition。在第一个版本中,我们对输入数组求反,然后使用argpartition来获得最小的k索引的索引,从而有效地得到最大的索引,而在第二个版本中,我们得到了第一个最小的索引,然后选择剩下的最大的索引。在

另外,这里值得一提的是,使用argpartition,我们没有按排序顺序得到索引。如果需要排序顺序,我们需要向np.argpartition提供范围数组,如^{}所述。在

样本运行-

1)二维盒:

^{pr2}$

2)3D案例:

In [46]: a # 3D array
Out[46]: 
array([[[20, 98, 27, 73],
        [33, 78, 48, 59],
        [28, 91, 64, 70]],

       [[47, 34, 51, 19],
        [73, 38, 63, 94],
        [95, 25, 93, 64]]])

In [47]: k_largest_index_argsort(a, k=2)
Out[47]: 
array([[0, 0, 1],
       [1, 2, 0]])

运行时测试-

In [56]: a = np.random.randint(0,99999999999999,(3000,4000))

In [57]: %timeit k_largest_index_argsort(a, k=10)
1 loops, best of 3: 2.18 s per loop

In [58]: %timeit k_largest_index_argpartition_v1(a, k=10)
10 loops, best of 3: 178 ms per loop

In [59]: %timeit k_largest_index_argpartition_v2(a, k=10)
10 loops, best of 3: 128 ms per loop

相关问题 更多 >