如何高效过滤Numpy有序唯一数组中间距过近的元素？

给定有序且元素唯一的Numpy数组，我们需要过滤掉所有与前一个保留元素间距≤指定值的元素。现有Python循环实现可以满足需求，但处理大规模数组时效率不足，下面提供更高效的优化方案。

需求示例 #

输入数组 #

import numpy as np
sample_array = np.array([1, 7, 15, 16, 18, 19, 26, 33])

不同间距参数的期望结果 #

• 当dist = 1时，移除与前一个保留元素间距≤1的16、19，结果：

  result = np.array([1, 7, 15, 18, 26, 33])
  

• 当dist = 3时，移除与前一个保留元素间距≤3的16、18，结果：

  result = np.array([1, 7, 15, 19, 26, 33])
  

现有循环实现 #

当前用Python循环+列表的实现逻辑清晰，但处理大数据量时速度较慢：

delta = dist  # delta对应参数dist
it = np.nditer(sample_array[1:])
result_list = [sample_array[0]]

for i in it:
    if (i - result_list[-1]) > delta:
        result_list.append(i)

result = np.array(result_list)

高效优化方案 #

方案1：Numpy掩码优化循环 #

用Numpy数组替代Python列表存储中间状态，减少列表append的开销，同时利用Numpy的数组操作提升效率：

def filter_close_elements(arr, dist):
    if arr.size == 0:
        return arr
    # 初始化掩码，标记哪些元素需要保留
    mask = np.ones(arr.shape, dtype=bool)
    last_kept = arr[0]
    for i in range(1, arr.size):
        if arr[i] - last_kept <= dist:
            mask[i] = False
        else:
            last_kept = arr[i]
    return arr[mask]

方案2：Numba编译加速 #

如果数组规模极大，用Numba将循环编译为机器码，能获得数量级的速度提升：

from numba import jit

@jit(nopython=True)
def filter_close_elements_numba(arr, dist):
    if arr.size == 0:
        return arr
    result = [arr[0]]
    last = arr[0]
    for num in arr[1:]:
        if num - last > dist:
            result.append(num)
            last = num
    return np.array(result)

方案说明 #

因为逻辑依赖前一个保留元素的状态，无法完全用纯矢量化操作实现（比如np.diff只能计算相邻元素差，不符合需求）。上面两种方案中：

• 方案1适合中等规模数组，代码无需额外依赖；
• 方案2适合超大规模数组，需要提前安装Numba（pip install numba）。

内容的提问来源于stack exchange，提问作者Julz