首页 > 解决方案 > numpy - 最小子数组

问题描述

我有一个 3 维 numpy ndarray

例如,一个 4x4x2 数组如下:

array = [
 [ [7 1] [8 0] [2 0] [7 1] ]
 [ [5 4] [1 4] [6 7] [8 1] ]
 [ [3 2] [4 5] [8 6] [6 2] ]
 [ [6 4] [1 2] [5 5] [7 1] ]
]

我需要找到一个最小嵌套的最内层数组(一个有两个数字的数组)和它的索引,像 Python 通常那样进行比较:比较第一对元素,如果相等,比较下一对,...

示例数组的预期结果:value=[1 2],indices=(3, 1)

用于查找元素本身的纯 Python 代码如下所示:

min(nested2 for nested1 in array for nested2 in nested1)

但是,我更喜欢 numpy 解决方案,因为数组非常庞大......

标签: numpymultidimensional-arraynumpy-ndarray

解决方案


您想要的行为取决于 Python 对列表列表的词法排序。也就是说,子列表被比较为:

In [275]: [1,2]<[2,0]                                                                          
Out[275]: True

np.sort只对结构化数组和复杂值进行词法排序。

In [288]: alist = [ 
     ...:  [ [7, 1], [8, 0], [2, 0], [7, 1] ], 
     ...:  [ [5, 4], [1, 4], [6, 7], [8, 1] ], 
     ...:  [ [3, 2], [4, 5], [8, 6], [6, 2] ], 
     ...:  [ [6, 4], [1, 2], [5, 5], [7, 1] ] 
     ...: ]                                                                                    
In [289]: arr = np.array(alist)                                                                
In [290]: arr                                                                                  
Out[290]: 
array([[[7, 1],
        [8, 0],
        [2, 0],
        [7, 1]],

       [[5, 4],
        [1, 4],
        [6, 7],
        [8, 1]],

       [[3, 2],
        [4, 5],
        [8, 6],
        [6, 2]],

       [[6, 4],
        [1, 2],
        [5, 5],
        [7, 1]]])

让我们尝试复杂的路线:

In [291]: x = np.dot(arr, [1,1j])                                                              
In [292]: x                                                                                    
Out[292]: 
array([[7.+1.j, 8.+0.j, 2.+0.j, 7.+1.j],
       [5.+4.j, 1.+4.j, 6.+7.j, 8.+1.j],
       [3.+2.j, 4.+5.j, 8.+6.j, 6.+2.j],
       [6.+4.j, 1.+2.j, 5.+5.j, 7.+1.j]])
In [293]: np.min(x)                                                                            
Out[293]: (1+2j)
In [294]: np.argmin(x)                                                                         
Out[294]: 13
In [295]: np.unravel_index(13, x.shape)                                                        
Out[295]: (3, 1)

一些时间测试:

In [302]: timeit min(nested2 for nested1 in alist for nested2 in nested1)                      
2.24 µs ± 19.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [303]: timeit min(arr.reshape(-1,2).tolist())                                               
2.53 µs ± 13.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
In [304]: timeit np.min(arr.dot([1,1j]))                                                       
19.5 µs ± 46.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

当阵列更大时,复杂的路线可能会更快,但对于这个样本,它是劣质的。

===

结构化数组方法:

In [320]: import numpy.lib.recfunctions as rf                                                  
In [321]: rf.unstructured_to_structured(arr, names=['x','y'])                                  
Out[321]: 
array([[(7, 1), (8, 0), (2, 0), (7, 1)],
       [(5, 4), (1, 4), (6, 7), (8, 1)],
       [(3, 2), (4, 5), (8, 6), (6, 2)],
       [(6, 4), (1, 2), (5, 5), (7, 1)]],
      dtype=[('x', '<i8'), ('y', '<i8')])
In [322]: np.argsort(rf.unstructured_to_structured(arr, names=['x','y']).ravel())              
Out[322]: array([13,  5,  2,  8,  9,  4, 14, 11, 12,  6,  0,  3, 15,  1,  7, 10])
In [323]: timeit np.argsort(rf.unstructured_to_structured(arr, names=['x','y']).ravel())       
41.3 µs ± 192 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

推荐阅读