在二维数组/矩阵中查找k个上限值的索引

时间:2019-07-18 23:47:37

标签: python numpy

我有一个带有值的2D矩阵,我想找到前5个值的索引。 例如

matrix([[0.17542851, 0.13199346, 0.01579704, 0.01429822, 0.01302919],
        [0.13279703, 0.12444886, 0.04742024, 0.03114371, 0.02623729],
        [0.13502306, 0.07815065, 0.07291175, 0.03690815, 0.02163695],
        [0.19032505, 0.15853737, 0.05889324, 0.02791679, 0.02699252],
        [0.1695696 , 0.14538635, 0.07127667, 0.04997876, 0.02580234]])

我想得到(0,3), (0,1), (0,4), (3,1), (4,1)

我搜索并尝试了许多解决方法,包括np.argmax(), np.argsort(), np.argpartition(),但均未获得良好的结果。 例如:

>>np.dstack(np.unravel_index(np.argsort(a.ravel(),axis=None), a.shape))

array([[[0, 4],
        [0, 3],
        [0, 2],
        [2, 4],
        [4, 4],
        [1, 4],
        [3, 4],
        [3, 3],
        [1, 3],
        [2, 3],
        [1, 2],
        [4, 3],
        [3, 2],
        [4, 2],
        [2, 2],
        [2, 1],
        [1, 1],
        [0, 1],
        [1, 0],
        [2, 0],
        [4, 1],
        [3, 1],
        [4, 0],
        [0, 0],
        [3, 0]]], dtype=int64)

此结果毫无意义。 请注意,我想要原始索引,我不在乎顺序(只是想以任何顺序排在前5位,尽管升序会更好)

4 个答案:

答案 0 :(得分:2)

您的样本:

n = np.array([[0.17542851, 0.13199346, 0.01579704, 0.01429822, 0.01302919],
        [0.13279703, 0.12444886, 0.04742024, 0.03114371, 0.02623729],
        [0.13502306, 0.07815065, 0.07291175, 0.03690815, 0.02163695],
        [0.19032505, 0.15853737, 0.05889324, 0.02791679, 0.02699252],
        [0.1695696 , 0.14538635, 0.07127667, 0.04997876, 0.02580234]])

您的输出不是前5个值的指标。前5个值是

array([0.14538635, 0.15853737, 0.1695696 , 0.17542851, 0.19032505])

要获取其索引:sort并使用isin标记其位置True。最后,使用argwhere获得他们的偏爱

np.argwhere(np.isin(n, np.sort(n, axis=None)[-5:]))

Out[324]:
array([[0, 0],
       [3, 0],
       [3, 1],
       [4, 0],
       [4, 1]], dtype=int32)

答案 1 :(得分:2)

np.argpartition应该是一个很好的工具(高效的工具),它可以在不保持顺序的情况下获得前k个索引。因此,对于数组数据a,它应该是-

In [43]: np.c_[np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)]
Out[43]: 
array([[4, 1],
       [3, 1],
       [4, 0],
       [0, 0],
       [3, 0]])

为说明起见,我们将其分解为单个处理步骤-

# Get partitioned indices such that the last 5 indices refer to the top 5
# values taken globally from the input array. Refer to docs for more info
# Note that it's global because we will flatten it. 
In [9]: np.argpartition(a.ravel(),-5)
Out[9]: 
array([14, 24,  2,  3,  4, 23, 22,  7,  8,  9, 19, 18, 17, 13, 12, 11,  6,
        1,  5, 10, 21, 16, 20,  0, 15])

# Get last 5 indices, which are the top 5 valued indices
In [10]: np.argpartition(a.ravel(),-5)[-5:]
Out[10]: array([21, 16, 20,  0, 15])

# Convert the global indices back to row-col format
In [11]: np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)
Out[11]: (array([4, 3, 4, 0, 3]), array([1, 1, 0, 0, 0]))

# Stack into two-columnar array
In [12]: np.c_[np.unravel_index(np.argpartition(a.ravel(),-5)[-5:],a.shape)]
Out[12]: 
array([[4, 1],
       [3, 1],
       [4, 0],
       [0, 0],
       [3, 0]])

对于a中的矩阵数据,应为-

In [48]: np.dstack(np.unravel_index(np.argpartition(a.ravel(),-5)[:,-5:],a.shape))
Out[48]: 
array([[[4, 1],
        [3, 1],
        [4, 0],
        [0, 0],
        [3, 0]]])

因此,与数组相比,唯一的区别在于使用np.dstack,因为对于矩阵数据,数据始终保持为2D。

请注意,这些是您从最后5行获得的结果。

答案 2 :(得分:1)

假设您有一个列表列表:

In [112]: M                                                                                                                                                                                                                                                                                                                   
Out[112]: 
[[0, 1, 2, 3, 4],
 [5, 6, 7, 8, 9],
 [10, 11, 12, 13, 14],
 [15, 16, 17, 18, 19],
 [20, 21, 22, 23, 24]]

In [113]: heapq.nlargest(5, ((r,c) for r in range(len(M)) for c in range(len(M[r]))), key=lambda t: M[t[0]][t[1]])                                                                                                                                                                                                            
Out[113]: [(4, 4), (4, 3), (4, 2), (4, 1), (4, 0)]

别忘了import heapq

答案 3 :(得分:0)

我是从引用@Divakar(非常优雅和快速)答案的 question 来到这里的。

排名的一个常见问题是如何处理重复项(并列)。

在某些情况下,需要使用“dense rank”,其中 [4, 7, 7, 9] 将排名(升序):[0, 1, 1, 2]

相比之下,@Divakar 的回答本质上反映了“ordinal ranking”,其中 [4, 7, 7, 9] 将排名(升序)[0, 1, 2, 3]。在“top k”问题中,这可能有点违反直觉。例如,在:

b = np.array(
    [[8, 6, 3],
    [6, 7, 2],
    [0, 8, 9]])

具有排名 k=2 和(并假设降序),它给出:

k = 2
>>> np.c_[np.unravel_index(np.argpartition(b.ravel(),-k)[-k:], b.shape)]
array([[2, 1],
       [2, 2]])

它对应于 9 并且只有一个 8 值,而忽略了另一个 8 值。

如果有人对“密集排名”感兴趣,我会提出以下建议(以“任何顺序”返回前 k 个值的所有索引——实际上,按索引顺序):

def topk_indices(a, k):
    _, rix = np.unique(-a, return_inverse=True)
    return np.c_[np.unravel_index(np.where(rix < k)[0], a.shape)]

在 OP 的数组上:

>>> topk_indices(a, 5)
array([[0, 0],
       [3, 0],
       [3, 1],
       [4, 0],
       [4, 1]])

在上面的 int 数组上:

>>> topk_indices(b, 2)
array([[0, 0],
       [2, 1],
       [2, 2]])

性能

在性能方面,@Divakar 的答案比这快 5 到 10 倍,适用于不同维度和参数的大量测试。所以如果你认为你没有关系,或者你不在乎,那就用它来代替这个。

举个例子:

a = np.random.randint(0, 10, (1_000_000, 2))
t0 = %timeit -o topk_indices(a, 5)
# 157 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

t1 = %timeit -o divakar_topk_indices(a, 5)
# 25.1 ms ± 49.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

>>> t0.average / t1.average
6.24

作为旁注,我必须对整个数组 (O(n log n)) 进行排序才能找到 top-k... 一个更合乎逻辑的 {{1 }} 方法扩展性更好 (heapq),但具有更大的常数乘数(仅 O(n log k) 需要 211 毫秒,这只是返回值,而不是索引。