在 Python 中查找二维数组中非零元素位置的快速方法

Question

我提出这个问题的原因是我认为 numpy.nonzero 函数没有得到应有的优化。下面的例子说明了这个事实：

a=np.random.random((1000,1000))
a[a<0.5]=0
timeit.timeit(lambda:np.nonzero(a),number=100)/100.0
#gives 0.014248089790344238 secs 
timeit.timeit(lambda:np.nonzero(a>0),number=100)/100.0
#gives 0.010561790466308594 secs

我不明白为什么会发生这种情况，因为我相信这两行对于每个输入都会产生相同的结果，无论 dtype 是什么，因此我将其命名为错误，如果错误，请纠正我。无论如何，严重缺乏优化。另一种提高查找非零元素速度的方法是，如果多次调用该函数，输入数组的大小是静态的且数组是稀疏的，则创建一个包含输入数组所有位置的 3d 数组，然后执行这个数组的布尔切片，使用相应的掩码，即：

def check_speed(a,positions,thres):
    a[a<thres]=0
    print 'Thresholding at:',thres
    print '\t Number of nonzero elements:',np.sum(a>0)
    print '\t Using numpy nonzero:', timeit.timeit(lambda:np.transpose(np.nonzero(a)),number=100)/100.0, 's'
    print '\t Using improved numpy nonzero:', timeit.timeit(lambda:np.transpose(np.nonzero(a>0)),number=100)/100.0, 's'
    print '\t Using readily made matrix:', timeit.timeit(lambda:positions[a>0],number=100)/100.0 ,'s'

a=np.random.random((1000,1000))
positions = np.transpose(np.nonzero(np.ones_like(a))).reshape(a.shape + (2,))
check_speed(a,positions,0.5)
check_speed(a,positions,0.8)
check_speed(a,positions,0.9)
check_speed(a,positions,0.995)
check_speed(a,positions,0.998)
check_speed(a,positions,0.999)
check_speed(a,positions,0.9995)

结果：

Thresholding at: 0.5
     Number of nonzero elements: 499248
     Using numpy nonzero: 0.0171903395653 s
     Using improved numpy nonzero: 0.0139024186134 s
     Using readily made matrix: 0.0163935399055 s    
Thresholding at: 0.8
     Number of nonzero elements: 199189
     Using numpy nonzero: 0.0108882308006 s
     Using improved numpy nonzero: 0.00774354934692 s
     Using readily made matrix: 0.00900489091873 s
Thresholding at: 0.9
     Number of nonzero elements: 99493
     Using numpy nonzero: 0.00896275043488 s
     Using improved numpy nonzero: 0.00562391996384 s
     Using readily made matrix: 0.00671306848526 s
Thresholding at: 0.995
     Number of nonzero elements: 4849
     Using numpy nonzero: 0.00732887983322 s
     Using improved numpy nonzero: 0.00353765010834 s
     Using readily made matrix: 0.00311932086945 s
Thresholding at: 0.998
     Number of nonzero elements: 1930
     Using numpy nonzero: 0.00729090929031 s
     Using improved numpy nonzero: 0.00361618041992 s
     Using readily made matrix: 0.00302290916443 s
Thresholding at: 0.999
     Number of nonzero elements: 959
     Using numpy nonzero: 0.00952008008957 s
     Using improved numpy nonzero: 0.00364511966705 s
     Using readily made matrix: 0.00303998947144 s
Thresholding at: 0.9995
     Number of nonzero elements: 490
     Using numpy nonzero: 0.00758473873138 s
     Using improved numpy nonzero: 0.00350986003876 s
     Using readily made matrix: 0.0028409409523 s

在 1000*1000 数组中低于 10000 个非零元素，该方法似乎可以提供更快的结果。我猜这是因为内存操作和计算之间存在瓶颈。没有理由研究 scipy 等价物，因为根据经验 scipy 相对于 numpy 模块非常慢。所以，我的问题是：

有谁知道在数组中查找非零元素位置的更好方法？预先感谢您的回答。

更新答案

与@Roxanne 的回答合作：

def check_speed(a,positions,thres):
    a[a<thres]=0
    print 'Thresholding at:',thres
    print '\t Number of nonzero elements:',np.sum(a>0)
    print '\t Using cv2:', timeit.timeit(lambda:np.fliplr(cv2.findNonZero((a>0).astype(np.uint8)).squeeze()),number=100)/100.0, 's'
    print '\t Using numpy nonzero:', timeit.timeit(lambda:np.transpose(np.array(np.nonzero(a))),number=100)/100.0, 's'
    print '\t Using improved numpy nonzero:', timeit.timeit(lambda:np.transpose(np.nonzero(a>0)),number=100)/100.0, 's'
    print '\t Using readily made matrix:', timeit.timeit(lambda:positions[a>0],number=100)/100.0 ,'s'
a=np.random.random((1000,1000))
positions = np.transpose(np.nonzero(np.ones_like(a))).reshape(a.shape + (2,))
a[a<0.5]=0
check_speed(a,positions,0.5)
check_speed(a,positions,0.8)
check_speed(a,positions,0.9)
check_speed(a,positions,0.995)
check_speed(a,positions,0.998)
check_speed(a,positions,0.999)
check_speed(a,positions,0.9995)

给出：

Thresholding at: 0.5
     Number of nonzero elements: 499806
     Using cv2: 0.00592091083527 s
     Using numpy nonzero: 0.0176041412354 s
     Using improved numpy nonzero: 0.0138215017319 s
     Using readily made matrix: 0.0165240502357 s
Thresholding at: 0.8
     Number of nonzero elements: 200022
     Using cv2: 0.00379456996918 s
     Using numpy nonzero: 0.0111160302162 s
     Using improved numpy nonzero: 0.00776970148087 s
     Using readily made matrix: 0.00905373096466 s
Thresholding at: 0.9
     Number of nonzero elements: 99835
     Using cv2: 0.00287639141083 s
     Using numpy nonzero: 0.00894243955612 s
     Using improved numpy nonzero: 0.00543779850006 s
     Using readily made matrix: 0.00669163942337 s
Thresholding at: 0.995
     Number of nonzero elements: 4931
     Using cv2: 0.00203591108322 s
     Using numpy nonzero: 0.00747500181198 s
     Using improved numpy nonzero: 0.00370697021484 s
     Using readily made matrix: 0.00331105947495 s
Thresholding at: 0.998
     Number of nonzero elements: 2018
     Using cv2: 0.00203846931458 s
     Using numpy nonzero: 0.00735512971878 s
     Using improved numpy nonzero: 0.0036293721199 s
     Using readily made matrix: 0.00307144880295 s
Thresholding at: 0.999
     Number of nonzero elements: 1038
     Using cv2: 0.00200258970261 s
     Using numpy nonzero: 0.00735764026642 s
     Using improved numpy nonzero: 0.00369818925858 s
     Using readily made matrix: 0.00302011013031 s
Thresholding at: 0.9995
     Number of nonzero elements: 521
     Using cv2: 0.0019414305687 s
     Using numpy nonzero: 0.00776562213898 s
     Using improved numpy nonzero: 0.00395655155182 s
     Using readily made matrix: 0.00321444988251 s

这是迄今为止我使用 Python 看到的最好的结果。

Answer 1

上次我检查cv2.findNonZero 比numpy.nonzero 快。但是你可能需要进行一些转换，它们是：np.fliplr(cv2.findNonZero((a>0).astype(np.uint8)).squeeze())