基于Numpy的高效实现：将连续非零元素求和替换至后续首个零位

现有Numpy数组：

import numpy as np
arr = np.array([[0, 0, 2, 5, 0, 0, 1, 8, 0, 3, 0],
                [1, 2, 0, 0, 0, 0, 5, 7, 0, 0, 0],
                [8, 5, 3, 9, 0, 1, 0, 0, 0, 0, 1]])

需要转换为目标数组：

[[0, 0, 0, 0, 7, 0, 0, 0, 9, 0, 3]
 [0, 0, 3, 0, 0, 0, 0, 0, 12, 0, 0]
 [0, 0, 0, 0, 25, 0, 1, 0, 0, 0, 0]]

处理规则：

• 逐行遍历，遇到0直接跳过
• 遇到非零元素时，累加后续连续的非零元素，直到碰到第一个0
• 将这个0替换为累加和，同时把之前参与累加的非零元素全部置0

当前已用循环实现，但处理大行数数组时效率极低，需要基于Numpy的向量化高效解决方案。

核心思路是利用Numpy向量化操作，通过标记连续非零块边界、计算块内和，再将和映射到对应位置，同时清空原非零块。

import numpy as np

def transform_array(arr):
    # 复制原数组避免修改输入
    result = arr.copy()
    rows, cols = arr.shape
    
    # 为每一行添加首尾0，处理行首/行尾非零的边界情况
    padded = np.pad(arr, ((0,0), (1,1)), mode='constant')
    
    # 标记非零元素位置，对比相邻元素找到块边界
    non_zero = padded != 0
    # 连续非零块的起始位置：当前非零，前一个为零
    starts = (non_zero[:, 1:-1] & ~non_zero[:, :-2])
    # 连续非零块对应的结束后第一个0的位置：当前为零，前一个非零
    ends = (~non_zero[:, 1:-1] & non_zero[:, :-2])
    
    # 计算每行累积和，快速获取任意区间的和
    cumsum = np.cumsum(padded, axis=1)
    
    # 逐行处理每个连续非零块
    for i in range(rows):
        row_starts = np.where(starts[i])[0]
        row_ends = np.where(ends[i])[0]
        
        # 计算每个非零块的累加和
        block_sums = cumsum[i, row_ends + 1] - cumsum[i, row_starts]
        
        # 将原非零块置0
        for s, e in zip(row_starts, row_ends):
            result[i, s:e] = 0
        # 将对应0位置替换为块和
        result[i, row_ends] = block_sums
    
    return result

# 测试示例
arr = np.array([[0, 0, 2, 5, 0, 0, 1, 8, 0, 3, 0],
                [1, 2, 0, 0, 0, 0, 5, 7, 0, 0, 0],
                [8, 5, 3, 9, 0, 1, 0, 0, 0, 0, 1]])
output = transform_array(arr)
print(output)

代码说明 #

1. 数组填充：给每行首尾加0，解决行首非零（如第二行开头的1、2）和行尾非零（如第三行末尾的1）的边界问题。
2. 边界标记：通过相邻元素对比，精准定位连续非零块的起始点和对应的结束后第一个0的位置。
3. 累积和计算：利用np.cumsum实现区间和的快速计算，避免逐元素累加的低效操作。
4. 批量更新：通过索引批量完成置零和赋值操作，替代逐元素循环，大幅提升处理大规模数组的效率。

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