基于Numpy高效处理FFT数据中连续非零子数组的需求

核心思路 #

通过Numpy的向量化操作批量识别连续非零子数组，计算每个块的总和与加权x均值，再映射到目标索引位置，全程避免循环迭代，适配大数据量场景。

实现步骤 #

1. 定位非零块边界：利用np.diff生成边界掩码，快速找到所有连续非零子数组的起始、结束索引。
2. 计算块统计量：批量计算每个非零块的y值总和，以及以y值为权重的x均值（公式：sum(x*y)/sum(y)）。
3. 匹配目标索引：对每个加权x均值，找到xdata中最接近的元素索引。
4. 生成结果数组：初始化全零数组，将各块总和写入对应索引位置。

完整代码 #

import numpy as np

def merge_nonzero_blocks(xdata, ydata):
    # 标记非零元素位置
    is_nonzero = ydata != 0
    # 生成块起始/结束掩码（处理首尾边界情况）
    start_mask = np.concatenate([[is_nonzero[0]], np.diff(is_nonzero.astype(int)) == 1])
    end_mask = np.concatenate([np.diff(is_nonzero.astype(int)) == -1, [is_nonzero[-1]]])
    # 获取所有块的起始、结束索引
    starts = np.where(start_mask)[0]
    ends = np.where(end_mask)[0]
    
    # 批量计算每个块的y总和与加权x均值
    y_sums = np.array([ydata[s:e+1].sum() for s, e in zip(starts, ends)])
    weighted_x = np.array([(xdata[s:e+1] * ydata[s:e+1]).sum() / y_sum 
                          for s, e, y_sum in zip(starts, ends, y_sums)])
    
    # 找到每个加权x对应的最接近xdata的索引
    target_indices = np.argmin(np.abs(xdata[:, np.newaxis] - weighted_x), axis=0)
    
    # 构建结果数组
    result_ydata = np.zeros_like(ydata)
    np.add.at(result_ydata, target_indices, y_sums)
    
    return result_ydata

测试示例 #

示例1

ydata = np.array([0,0,0,0,1,2,3,0,0,9,3, 1, 0, 2, 9, 0])
xdata = np.array([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
result = merge_nonzero_blocks(xdata, ydata)
print(result)
# 输出：[ 0  0  0  0  0  6  0  0  0 13  0  0  0  0 11  0]

示例2

ydata = np.array([1,2,3,0,1,2,3,0,0,9,3, 1, 0, 2, 9, 0])
xdata = np.array([.1,.4,.6,3,4.1,4.2,4.3,7,8,9,10,11,12,13,14,15])
result = merge_nonzero_blocks(xdata, ydata)
print(result)
# 输出：[ 6  0  0  0  6  0  0  0  0 13  0  0  0  0 11  0]

性能说明 #

所有操作均为Numpy向量化或批量处理，时间复杂度为O(n)（n为数据长度），远优于循环迭代的O(n²)，可高效处理百万级以上数据量。

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