Python 沿已排序二维numpy数组的轴查找第一个非零值

我试图找到最快的方法，为二维排序数组的每一行找到第一个非零值。从技术上讲，数组中唯一的值是0和1，并且是“排序”的

例如，数组可能如下所示：

五=

我可以使用argmax函数

argmax(v, axis=1))

当它从0变为1时，我想这会对每一行进行彻底的搜索。我的阵列大小将合理（~2000x2000）。对于for循环中的每一行，argmax是否仍然比只使用searchsorted方法表现更好，还是有更好的替代方法

此外，数组将始终使一行的第一个位置始终>=其上方行中一的第一个位置（但不能保证最后几行中会有一个）。我可以利用for循环和一个“起始索引值”来利用这一点，每行的值等于前一行中第一个1的位置，但是我认为numpy argmax函数的性能仍然优于用python编写的循环，这是正确的

我只想对备选方案进行基准测试，但数组的边长可能会有很大的变化（从250到10000）。

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。

使用起来相当快：

它提供的元组的坐标值大于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...

如果您可以处理对于所有零的行不使用

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中使用axis的版本（如您评论中所建议的）：

对于速度比较（在您的示例阵列上），我得到以下结果：

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

---

我非常希望argmax函数更快。如果它是性能关键型的，您可以尝试用C语言编写一个扩展，argmax的好处是您可以指定一个轴，即

argmax（a，axis=1）

，它将使用C语言编写的循环遍历行，因此您不必使用python for循环，这应该会更慢。@user1554752：是的，但是如果您使用

argmax（a，轴=1）

那么

a

中的行之间存在歧义，这些行是

[1，x，x，x，]

或

[0,0,0]

，因为

argmax（a，轴=1）

对于任何一种情况都将返回

0

。您仍然需要在argmax返回的数组上循环以测试这种不确定性，不是吗？这就是我可以利用数据中的模式，其中第一个1永远不会位于其上一行中第一个1的左侧。一旦我从argmax获得数组（称为indx），我可以在它上面运行argmin。如果它返回一个值p！=0，那么从p向下的所有行都由零组成。

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

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