在 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

我不明白为什么会发生这种情况，因为我相信两条线都会为每个输入产生相同的结果，无论数据类型如何，因此我将其命名为错误，如果错误请纠正我。无论如何，严重缺乏优化。另一种提高查找非零元素速度的方法是，如果函数被多次调用，输入数组的大小是静态的并且数组是稀疏的，创建一个包含输入数组所有位置的 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())

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

Fast way to find nonzero elements positions in 2d array in Python

python

arrays

performance

2d

numpy

更新答案