python - 第一组非零值(通过忽略零的单次出现)

Question

这是我打算用 Python 做的:

我有一个数组 (freq_arr)。我想找到第一组非零元素的索引。我从头开始寻找非零元素，当我找到第一个非零元素(第一个元素是5，在下面的例子中)，我记录它的索引(4，在下面的例子中)。我搜索下一个，并记录其索引(将为 5)。如果遇到单个零，我想忽略它并继续搜索非零值。这样，我考虑值 5、6、0、8、9、0、1，索引为 4、5、6、7、8、9 和 10。在这些值之后，有五个零，因此我停止搜索.输出中最多可以存在两个零，然后继续搜索。但是，如果我遇到 3 个或更多个零，我想停止搜索。

输入:

freq_arr = np.array([0, 0, 0, 0, 5, 6, 0, 8, 9, 0, 1, 0, 0, 0, 0, 3, 6, 0])

输出:

out_arr_indices = [4, 5, 6, 7, 8, 9, 10]

我知道使用 for 循环对此进行编码，但我想避免使用它，因为它效率不高。请让我知道如何做到这一点。

数组将是单维的。每个元素将在 5000 到 20000 的范围内。

Answer 1

这是一种使用切片和 argmax 的方法(检测非零和零)-

def start_stop_indices(freq_arr, W=3):
    nnz_mask = freq_arr!=0
    start_idx = nnz_mask.argmax()
    m0 = nnz_mask[start_idx:]
    kernel = np.ones(W,dtype=int)
    last_idx = np.convolve(m0, kernel).argmin() + start_idx - W
    return start_idx, last_idx

样本运行-

In [203]: freq_arr
Out[203]: array([0, 0, 0, 0, 5, 6, 0, 8, 9, 0, 1, 0, 0, 0, 0, 3, 6, 0])

In [204]: start_stop_indices(freq_arr, W=3)
Out[204]: (4, 10)

In [205]: start_stop_indices(freq_arr, W=2)
Out[205]: (4, 10)

In [206]: start_stop_indices(freq_arr, W=1)
Out[206]: (4, 5)

这里还有一个针对length = 3的固定窗口搜索，避免使用卷积，更多的使用切片 -

def start_stop_indices_v2(freq_arr):
    nnz_mask = freq_arr!=0
    start_idx = nnz_mask.argmax()
    m0 = nnz_mask[start_idx:]
    idx0 = (m0[:-2] | m0[1:-1] | m0[2:]).argmin()
    last_idx = idx0 + start_idx - 1
    return  start_idx, last_idx