python - 沿排序的二维 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 循环中的每一行执行搜索排序方法，还是有更好的选择？

此外，数组总是这样，一行中 one 的第一个位置总是 >= 其上一行中 one 的第一个位置(但不保证在最后几行)。我可以用一个 for 循环和一个“起始索引值”来利用它，每一行等于前一行第一个 1 的位置，但是我认为 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.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 中使用轴的版本(如您的评论中所建议):

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