查找numpy数组中的第一个和最后一个匹配项

i2loujxw  于 2022-12-23  发布在  其他
关注(0)|答案(3)|浏览(183)

我有一个(大的)numpy布尔值数组,其中True表示另一个先前处理过的信号的过零情况。在我的例子中,我只对数组中第一个和最后一个True值的索引感兴趣,因此我宁愿不使用np.wherenp.argwhere,以避免不必要地遍历整个数组。
我不想把我的数组转换成list,而是使用index()方法。
到目前为止,我想出的最好的解决方案如下:

# Alias for compactness
enum = enumerate

# Example array, desired output = (3, -5))
a = np.array([False, False, False, True, ..., True, False, False, False, False])

out = next(idx for idx in zip((i for i, x in enum(a) if x), (-i-1 for i, x in enum(a[::-1]) if x)))

print(out)
# (3, -5)

但我还是觉得它有点笨重,想知道Numpy是否已经用另一种语法实现了这样的功能,是吗?

jtjikinw

jtjikinw1#

使用numpy的nonzero函数,该函数返回输入数组中所有非零元素的索引:

a = np.array([False, False, False, True, ..., True, False, False])

# Find the indices of the non-zero elements
indices = np.nonzero(a)

# The first and last indices will be the first and last True values
first = indices[0][0]
last = indices[0][-1]
6rqinv9w

6rqinv9w2#

这个怎么样?

import numpy as np

a = np.array([False, False, False, True, False, True, False, False, False, False])

out = (np.argmax(a), a.size - np.argmax(a[::-1]) - 1)

np.argmax()在第一个True值处停止,因此这可能会解决您不想迭代整个数组的问题?

lf5gs5x2

lf5gs5x23#

基准测试结果

因此,我决定在一个典型案例中比较所提出的解决方案。这三种方法是:

@timeit
def gener_sol(a):
    # My OP proposal
    return next(idx for idx in zip((i for i, x in enumerate(a) if x), (a.size-i-1 for i, x in enumerate(a[::-1]) if x)))

@timeit
def nonzero_sol(a):
    # kalzso's proposal
    return np.nonzero(a)[0][[0, -1]]

@timeit
def argmax_sol(a):
    # jan992's proposal
    return np.argmax(a), a.size - np.argmax(a[::-1]) - 1

注意:@timeit是一个定制的 Package 器,我必须测量函数的执行时间。除了它是否完美的问题外,它对所有函数都是一样的,所以它应该是一个公平的比较。
下面是在循环中执行它们的代码:

for rep in range(10000):
    N = 2**20  # Array size
    depth = N // 10  # max depth from the edges for first and last True
    
    arr = np.random.randint(0, 2, N).astype(bool)

    # First and last True values are randomly located at both ends
    first = random.randint(0, depth)
    arr[:first] = False
    arr[first]  = True

    last = random.randint(0, depth)
    arr[-last:] = False
    arr[-last]  = True

    arg_res = argmax_sol(arr)
    non_res = nonzero_sol(arr)
    gen_res = gener_sol(arr)

    # Ensure they all return the same
    assert arg_res == tuple(non_res) == gen_res

研究结果如下:

    • N = 2^20**(10000次运行)
argmax_sol....:   10000 runs,       6.38 s  (638.01 μs/run)
nonzero_sol...:   10000 runs,     13.124 s  (1.312 ms/run)
gener_sol.....:   10000 runs,     42.695 s  (4.27 ms/run)
    • N = 2^25**(100次运行)
argmax_sol....:     100 runs,      1.891 s  (18.908 ms/run)
nonzero_sol...:     100 runs,      4.125 s  (41.25 ms/run)
gener_sol.....:     100 runs,     13.799 s  (137.99 ms/run)

因此,在这个场景中,argmax解决方案显然是赢家,而generator显然是输家。然而,我注意到第一个和最后一个索引所在的depth(在这个例子中最多为N/10)可能会产生重大影响。在我的例子中,合理的值是depth=10000(甚至更小)。因此,如果我重复这个实验,我会得到:

    • N = 2^20**(10000次运行)depth = 10000
gener_sol.....:     100 runs,    39.967 ms  (399.671 μs/run)
argmax_sol....:     100 runs,    58.851 ms  (588.505 μs/run)
nonzero_sol...:     100 runs,   138.077 ms  (1.381 ms/run)

如您所见,在本例中,generator产生了更好的结果,如果depth可以绑定到更小的数,则改进会更大。
例如,参见:

    • N = 2^20**(10000次运行)depth=1000
gener_sol.....:   10000 runs,      0.947 s  (94.753 μs/run)
argmax_sol....:   10000 runs,      6.082 s  (608.196 μs/run)
nonzero_sol...:   10000 runs,     14.282 s  (1.428 ms/run)
    • N = 2^25**(1000次运行)depth=1000
gener_sol.....:    1000 runs,      0.053 s  (53.08 μs/run)
argmax_sol....:    1000 runs,      18.84 s  (18.84 ms/run)
nonzero_sol...:    1000 runs,     44.058 s  (44.058 ms/run)

哇,这是一个巨大的改进!当N = 2^20和2^25时,分别比numpy的argmax快了~6倍和~350倍。因此,在适当的情况下,不应该低估for循环/生成器的能力。
最后,为了给kalszo的nonzero解决方案留下一句好话,我确实发现了它优于其他解决方案的一种情况,那就是除了两个位置(要查找的位置)之外,数组到处都是False
因此,更改行:

# arr = np.random.randint(0, 2, N).astype(bool)
arr = np.zeros(N).astype(bool)
    • N = 2^20**depth=N//10,只有两个真值
nonzero_sol...:   10000 runs,      1.612 s  (161.186 μs/run)
argmax_sol....:   10000 runs,      6.482 s  (648.178 μs/run)
gener_sol.....:   10000 runs,     44.625 s  (4.463 ms/run)

并且具有较浅的深度:

    • N = 2^20**(10000次运行)depth=10000,只有两个真值
nonzero_sol...:   10000 runs,      1.571 s  (157.093 μs/run)
gener_sol.....:   10000 runs,      4.208 s  (420.791 μs/run)
argmax_sol....:   10000 runs,      6.235 s  (623.512 μs/run)
    • N = 2^25**(1000次运行)depth=10000,只有两个真值
gener_sol.....:    1000 runs,      0.433 s  (433.125 μs/run)
nonzero_sol...:    1000 runs,      5.427 s  (5.427 ms/run)
argmax_sol....:    1000 runs,     19.128 s  (19.128 ms/run)
    • N = 2^20**(10000次运行)depth=1000,只有两个真值
gener_sol.....:   10000 runs,      0.962 s  (96.286 μs/run)
nonzero_sol...:   10000 runs,      1.637 s  (163.687 μs/run)
argmax_sol....:   10000 runs,      6.522 s  (652.228 μs/run)

相对于阵列大小的小深度再次是有利于生成器表达式的关键因素。
总之,正确的解决方案在很大程度上取决于场景。对于我的特定情况,其中:

  • 数组很大(2^20 - 2^25)
  • 数组在未知位置具有多个True值
  • 第一个和最后一个True值相对靠近边

生成器表达式具有最好的性能,具有高达一个或两个数量级。

相关问题