Numpy按组查找首尾非零元素并填充区间的无循环实现方法

需求说明 #

• 输入两个等长Numpy数组：分组数组group、数值数组array
• 需求1：对每个分组，定位组内第一个、最后一个非零元素的索引
• 需求2：将每个分组首尾非零索引区间内的所有元素，替换为该组第一个非零元素的值
• 约束：不拆分数组、不使用Python层迭代遍历

全向量化实现代码 #

import numpy as np

def func4(group, arr):
    # 标记非零元素位置
    nonzero_mask = arr != 0
    # 计算每个分组的起始索引
    group_start = np.r_[0, np.where(np.diff(group) != 0)[0] + 1]
    group_length = np.diff(np.r_[group_start, len(group)])
    # 计算每个元素在所属组内的序号
    group_idx = np.arange(len(group)) - np.repeat(group_start, group_length)
    
    # 1. 计算每个分组首个非零的全局索引
    non_zero_group_idx = np.where(nonzero_mask, group_idx, len(group) + 1)
    first_nonzero_offset = np.minimum.reduceat(non_zero_group_idx, group_start)
    first_nonzero_global = first_nonzero_offset + group_start
    
    # 2. 计算每个分组最后一个非零的全局索引
    max_group_idx = np.maximum.reduceat(group_idx, group_start)
    rev_group_idx = max_group_idx[group] - group_idx
    rev_non_zero_group_idx = np.where(nonzero_mask, rev_group_idx, len(group) + 1)
    last_nonzero_offset = np.minimum.reduceat(rev_non_zero_group_idx, group_start)
    last_nonzero_global = group_start + max_group_idx - last_nonzero_offset
    
    # 3. 生成替换掩码：元素是否在所属分组的首尾非零区间内
    first_arr = np.repeat(first_nonzero_global, group_length)
    last_arr = np.repeat(last_nonzero_global, group_length)
    idx_arr = np.arange(len(arr))
    replace_mask = (idx_arr >= first_arr) & (idx_arr <= last_arr)
    
    # 4. 生成替换值数组
    first_val = arr[first_nonzero_global]
    replace_val_arr = np.repeat(first_val, group_length)
    
    # 5. 执行替换输出结果
    res = arr.copy()
    res[replace_mask] = replace_val_arr[replace_mask]
    # 如需返回各分组首尾非零索引，可新增返回值：dict(zip(np.unique(group), zip(first_nonzero_global, last_nonzero_global)))
    return res

验证示例 #

# 测试输入
group = np.array([0,0,0,0,1,1,1,1,1,1,2,2,2,2])
arr = np.array([1,2,3,0,0,2,0,3,4,0,0,0,0,5])
print(func4(group, arr))
# 输出：[1 1 1 0 0 2 2 2 2 0 0 0 0 5]，完全符合预期

性能说明 #

该实现全程使用Numpy原生向量化操作，无Python层循环、无数组拆分操作，3万长度数组运行100次仅耗时0.03秒，性能远高于常规遍历方案。

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