沿排序的二維numpy數組的軸找到第一個非零值

Question

我正在嘗試找到最快的方法來為二維排序數組的每一行找到第一個非零值。 從技術上講，數組中唯一的值是0和1，並且它是“已排序”的。

例如，數組可能如下所示：

v =

0 0 0 1 1 1 1 
0 0 0 1 1 1 1 
0 0 0 0 1 1 1 
0 0 0 0 0 0 1 
0 0 0 0 0 0 1 
0 0 0 0 0 0 1 
0 0 0 0 0 0 0

我可以使用argmax功能

argmax(v, axis=1))

找到它從零變為一，但我相信這將在每一行進行詳盡的搜索。 我的陣列大小合適（~2000x2000）。 argmax是否仍然優於僅針對for循環中的每一行執行搜索排序方法，還是有更好的選擇？

此外，數組總是這樣，一行的第一個位置總是> =它上面一行中的第一個位置（但不保證最后幾行中會有一個） ）。 我可以利用for循環和每個行的“起始索引值”來利用它，它等於前一行中前一個1的位置，但我認為numpy argmax函數仍然勝過用python編寫的循環。

我只是對備選方案進行基准測試，但陣列的邊長可能會發生很大變化（從250到10,000）。

Answer 1

argmax（）使用C級循環，它比Python循環快得多，所以我想即使你在Python中編寫一個智能算法，也很難擊敗argmax（），你可以使用Cython來加速：

@cython.boundscheck(False)
@cython.wraparound(False) 
def find(int[:,:] a):
    cdef int h = a.shape[0]
    cdef int w = a.shape[1]
    cdef int i, j
    cdef int idx = 0
    cdef list r = []
    for i in range(h):
        for j in range(idx, w):
            if a[i, j] == 1:
                idx = j
                r.append(idx)
                break
        else:
            r.append(-1)
    return r

在我的PC上用於2000x2000矩陣，它是100us vs 3ms。

Answer 2

使用np.where的速度相當快：

>>> a
array([[0, 0, 0, 1, 1, 1, 1],
       [0, 0, 0, 1, 1, 1, 1],
       [0, 0, 0, 0, 1, 1, 1],
       [0, 0, 0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0, 0, 0]])
>>> np.where(a>0)
(array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 4, 5]), array([3, 4, 5, 6, 3, 4, 5, 6, 4, 5, 6, 6, 6, 6]))

這會將元組傳遞給大於0的坐標值。

您還可以使用np.where來測試每個子數組：

def first_true1(a):
    """ return a dict of row: index with value in row > 0 """
    di={}
    for i in range(len(a)):
        idx=np.where(a[i]>0)
        try:
            di[i]=idx[0][0]
        except IndexError:
            di[i]=None    

    return di       

打印：

{0: 3, 1: 3, 2: 4, 3: 6, 4: 6, 5: 6, 6: None}

即，第0行：索引3> 0; 第4行：索引4> 0; 第6行：沒有大於0的索引

如您所料，argmax可能更快：

def first_true2():
    di={}
    for i in range(len(a)):
        idx=np.argmax(a[i])
        if idx>0:
            di[i]=idx
        else:
            di[i]=None    

    return di       
    # same dict is returned...

如果你可以處理所有naughts行None的邏輯，那么這仍然更快：

def first_true3():
    di={}
    for i, j in zip(*np.where(a>0)):
        if i in di:
            continue
        else:
            di[i]=j

    return di      

這是一個在argmax中使用軸的版本（如評論中所示）：

def first_true4():
    di={}
    for i, ele in enumerate(np.argmax(a,axis=1)):
        if ele==0 and a[i][0]==0:
            di[i]=None
        else:
            di[i]=ele

    return di          

對於速度比較（在您的示例數組上），我得到：

            rate/sec usec/pass first_true1 first_true2 first_true3 first_true4
first_true1   23,818    41.986          --      -34.5%      -63.1%      -70.0%
first_true2   36,377    27.490       52.7%          --      -43.6%      -54.1%
first_true3   64,528    15.497      170.9%       77.4%          --      -18.6%
first_true4   79,287    12.612      232.9%      118.0%       22.9%          --

如果我將其擴展到2000 X 2000 np數組，這是我得到的：

            rate/sec  usec/pass first_true3 first_true1 first_true2 first_true4
first_true3        3 354380.107          --       -0.3%      -74.7%      -87.8%
first_true1        3 353327.084        0.3%          --      -74.6%      -87.7%
first_true2       11  89754.200      294.8%      293.7%          --      -51.7%
first_true4       23  43306.494      718.3%      715.9%      107.3%          --