查找未被零包围的元素的局部最大值索引

Question

我正在尝试识别未被一维 numpy 数组的零包围的局部最大值的索引。

原始代码是：

max_idx = [
    i for i in range(1, len(elem_array) - 1)
    if ((elem_array[i - 1] < elem_array[i]) and (elem_array[i + 1] <= elem_array[i]))
    and ((elem_array[i - 1] != 0) or (elem_array[i + 1] != 0))
]

使用此代码使用数组：

elem_array = np.array([23,  0, 45,  0, 12, 13, 14,  0,  0,  0,  1, 67,  1])

结果是： max_idx = [6, 11] 。

重要提示：元素i可以大于或等于元素i+1 ，但仅大于元素i-1并且0只能在元素i的一侧，这就是为什么45不被识别为本地的原因最大。

我试图用scipy.signal.argrelextrema修改它，但这给了我结果： max_idx = [2, 6, 11] ，其中包含一个额外的元素。

并使用数组：

elem_array = np.array([0.0, 0.0, 0.0, 0.0, 0.07, 0.2, 0.4, 0.6, 0.8, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0])

结果是一个空数组，它应该是： max_idx = [10] 。

您对如何修改原始代码有什么建议吗？ 谢谢

Answer 1

像这样的循环很容易矢量化：

mask = (
    (elem_array[:-2] < elem_array[1:-1])
    & (elem_array[2:] <= elem_array[1:-1])
    & ((elem_array[:-2] != 0) | (elem_array[2:] != 0))
)
max_idx = np.nonzero(mask)[0] + 1

Answer 2

您可以使用numpy.lib.stride_tricks.sliding_window_view创建形状为 3 的滑动窗口，然后以矢量化方式应用条件：

import numpy as np

def get_local_maxima(a: np.array, window_shape: int = 3) -> np.array:
    
    mid_index = window_shape//2
    # Adding initial and final zeros and create the windo of given size
    window = np.lib.stride_tricks.sliding_window_view(np.array([0]*mid_index + [*a] + [0]*mid_index), window_shape)
    
    c1 = np.argmax(window, axis=1)==mid_index # first condition is that the max must be in the center of the window
    c2 = (window[:, [i for i in range(window_shape) if i!=mid_index]]!=0).any(axis=1) # second condition is that one among 0-th and 2-nd element must be non-zero
    
    return np.argwhere(c1 & c2)

a = np.array([23,  0, 45,  0, 12, 13, 14,  0,  0,  0,  1, 67,  1])
b = np.array([0.0, 0.0, 0.0, 0.0, 0.07, 0.2, 0.4, 0.6, 0.8, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0])

get_local_maxima(a)
array([[ 6],
       [11]])

get_local_maxima(b)
array([[10]])