沿着排序的二维numpy数组的轴寻找第一个非零值

Question

我试图找到找到两维排序数组的每一行的第一个非零值的最快方法。从技术上讲，数组中唯一的值是零和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）。对于for循环中的每一行，argmax仍然会超越搜索排序方法，还是有更好的选择？

另外，数组总是这样，一个行的第一个位置总是> =它上面一行的第一个位置（但不能保证会有一个位于最后几行）。我可以利用for循环和每行的“起始索引值”等于前一行第一个位置的位置，但是我正确地认为numpy argmax函数仍然会超出用python编写的循环。

我只是基准的替代品，但阵列的边缘长度可能会发生相当大的变化（从250到10,000）。

Answer 1

这是相当快的使用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更大的价值坐标。其中，以测试每个子阵列：

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%   -- 

Answer 2

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 

在我的电脑上的2000x2000矩阵，它是100us vs 3ms。