python之使用 np.argpartition 对多维数组中的值进行索引

我有一个这样的数组:

>>> a = np.arange(60).reshape([3,4,5]) 
>>> a 
array([[[ 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], 
        [25, 26, 27, 28, 29], 
        [30, 31, 32, 33, 34], 
        [35, 36, 37, 38, 39]], 
 
       [[40, 41, 42, 43, 44], 
        [45, 46, 47, 48, 49], 
        [50, 51, 52, 53, 54], 
        [55, 56, 57, 58, 59]]]) 

我想检索其中一个维度的前 k 个值。例如，我将选择 k=2 并沿着中间维度。

我试过使用 argpartition 并且它似乎做了正确的事情，但是我在使用它的输出从原始数组中检索值时遇到了问题。这是我使用 argpartition 的方式:

>>> indices = np.argpartition(a, 2, axis=1) 
>>> indices 
array([[[0, 0, 0, 0, 0], 
        [1, 1, 1, 1, 1], 
        [2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[0, 0, 0, 0, 0], 
        [1, 1, 1, 1, 1], 
        [2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[0, 0, 0, 0, 0], 
        [1, 1, 1, 1, 1], 
        [2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]]]) 
 
>>> indices[:,-2:,:] 
array([[[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]]]) 

但我无法通过使用这些索引进行切片来获取值。

>>> a[:,indices[:,-2:,:],:].shape 
(3, 3, 2, 5, 5) 

我期待看到形状为 (3,2,5) 的数组(因为我正在寻找沿中轴的前 2 个)，我想它看起来像这样:

>>> magic_output 
array([[[10, 11, 12, 13, 14], 
        [15, 16, 17, 18, 19]], 
 
       [[30, 31, 32, 33, 34], 
        [35, 36, 37, 38, 39]], 
 
       [[50, 51, 52, 53, 54], 
        [55, 56, 57, 58, 59]]]) 

如何使用 argpartition 中的索引访问值？

请您参考如下方法:

np.argpartition 得到最小的 k 索引。因此，要获得前 k 索引，我们需要沿所需轴使用否定输入数组。然后，我们需要使用这些索引通过 NumPy 的高级索引 索引到该轴并获得所需的输出。

因此，实现将是 -

k = 2 
m,n = a.shape[0], a.shape[2] 
idx = np.argpartition(-a,k,axis=1)[:,k-1::-1] 
out = a[np.arange(m)[:,None,None], idx, np.arange(n)] 

sample 运行-

1)输入:

In [180]: a 
Out[180]:  
array([[[ 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], 
        [25, 26, 27, 28, 29], 
        [30, 31, 32, 33, 34], 
        [35, 36, 37, 38, 39]], 
 
       [[40, 41, 42, 43, 44], 
        [45, 46, 47, 48, 49], 
        [50, 51, 52, 53, 54], 
        [55, 56, 57, 58, 59]]]) 

2) 提议的代码:

In [206]: k = 2 
     ...: m,n = a.shape[0], a.shape[2] 
     ...: idx = np.argpartition(-a,k,axis=1)[:,k-1::-1] 
     ...: out = a[np.arange(m)[:,None,None], idx, np.arange(n)] 
     ...:  

3) 检查中间结果和输出:

In [207]: idx 
Out[207]:  
array([[[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]], 
 
       [[2, 2, 2, 2, 2], 
        [3, 3, 3, 3, 3]]]) 
 
In [208]: out 
Out[208]:  
array([[[10, 11, 12, 13, 14], 
        [15, 16, 17, 18, 19]], 
 
       [[30, 31, 32, 33, 34], 
        [35, 36, 37, 38, 39]], 
 
       [[50, 51, 52, 53, 54], 
        [55, 56, 57, 58, 59]]])