如何对二维数组中标记的不规则切片逐行分段累加求和？

问题背景 #

你需要对大型二维数组中被掩码标记的连续分段进行独立累加——每个分段的累加从该段的第一个元素开始，不包含前一段的计算结果，这确实是np.cumsum无法直接实现的，因为它是全局累积的。

给定的输入数组和掩码：

import numpy as np

# 输入数组
a = np.array([
    [ 0., 1., 2., 3., 4., 5. , 6. , 7.],
    [ 6., 7., 8., 9., 10., 4.2, 4.3, 11.],
    [ 12., 13., 14., 15., 16., 4.2, 4.3, 17.],
    [ 18., 19., 20., 21., 22., 4.2, 4.3, 23.]
])

# 掩码（True表示需要参与分段累加的元素）
mask = np.array([
    [False, True, True, True, False, True, True , False],
    [False, False, False, True, True, True, True , False],
    [False, False, True, True, False, False, False, False],
    [False, True, True, False, False, False, True , True ]
])

预期输出：

array([[ 0., 1., 3., 6., 0., 5. , 11. , 0.],
       [ 0., 0., 0., 9., 19., 23.2, 27.5, 0.],
       [ 0., 0., 14., 29., 0., 0, 0, 0.],
       [ 0., 19., 39., 0., 0., 0, 4.3, 27.3]])

核心思路 #

要实现分段独立累加，关键要解决两个问题：

• 识别掩码中每一行的连续True分段的长度（即每个独立累加段的元素个数）
• 对每个分段内的元素做局部cumsum，分段之间重置累加起点

优化后的实现方案 #

这里基于你提供的代码进行梳理和优化，提升可读性和效率：

import numpy as np

def get_segment_sizes(mask, axis=1):
    """获取掩码中每个连续True分段的长度"""
    if mask.ndim != 2:
        raise ValueError("仅支持二维掩码")
    
    if axis == 1:
        # 按行处理，给每行前后添加False边界
        padded_mask = np.hstack([np.zeros((mask.shape[0],1), dtype=bool), mask, np.zeros((mask.shape[0],1), dtype=bool)])
        # 通过差分找到分段的起始（1）和结束（-1）位置
        diff = np.diff(padded_mask.astype(int), axis=1)
        starts = np.where(diff == 1)[1]
        ends = np.where(diff == -1)[1]
        sizes = ends - starts
    else:
        # 按列处理，给每列上下添加False边界
        padded_mask = np.vstack([np.zeros((1,mask.shape[1]), dtype=bool), mask, np.zeros((1,mask.shape[1]), dtype=bool)])
        diff = np.diff(padded_mask.astype(int), axis=0)
        starts = np.where(diff == 1)[0]
        ends = np.where(diff == -1)[0]
        sizes = ends - starts
    return sizes

def intervaled_cumsum(arr, segment_sizes):
    """对分段后的数组做局部独立累加"""
    arr_cumsum = arr.cumsum()
    # 计算每个分段的起始索引
    segment_starts = np.cumsum(segment_sizes) - segment_sizes
    # 生成偏移量：每个分段的第一个元素偏移为0，后续分段减去前一段的累加终点
    offsets = np.zeros_like(arr_cumsum)
    offsets[segment_starts[1:]] = arr_cumsum[segment_starts[1:]-1]
    # 减去累积的偏移量，得到局部累加结果
    local_cumsum = arr_cumsum - offsets.cumsum()
    return local_cumsum

def segmented_cumsum(arr, mask, axis=1):
    """对二维数组按掩码分段做独立累加"""
    if arr.shape != mask.shape:
        raise ValueError("数组和掩码形状必须一致")
    
    result = np.zeros_like(arr)
    if axis == 1:
        # 按行逐个处理
        for i in range(arr.shape[0]):
            row_data = arr[i]
            row_mask = mask[i]
            if not np.any(row_mask):
                continue
            # 获取当前行的分段长度
            sizes = get_segment_sizes(row_mask.reshape(1,-1), axis=1)
            # 提取该行需要累加的元素
            masked_data = row_data[row_mask]
            # 计算局部累加
            local_sum = intervaled_cumsum(masked_data, sizes)
            # 将结果放回原位置
            result[i][row_mask] = local_sum
    else:
        # 按列逐个处理
        for j in range(arr.shape[1]):
            col_data = arr[:,j]
            col_mask = mask[:,j]
            if not np.any(col_mask):
                continue
            sizes = get_segment_sizes(col_mask.reshape(-1,1), axis=0)
            masked_data = col_data[col_mask]
            local_sum = intervaled_cumsum(masked_data, sizes)
            result[:,j][col_mask] = local_sum
    return result

# 测试运行
output = segmented_cumsum(a, mask, axis=1)
print(output)

代码解释 #

1. get_segment_sizes：通过在掩码两端添加False边界，计算差分找到每个连续True分段的起始和结束位置，从而得到分段长度。
2. intervaled_cumsum：先对所有标记元素做全局cumsum，然后通过偏移量减去前一个分段的累加终点，实现每个分段的局部独立累加。
3. segmented_cumsum：遍历每一行（或列），提取标记元素，计算局部累加后放回原数组对应位置，未标记位置保持0。

性能优化建议 #

如果处理超大型数组，建议：

• 避免逐行/列遍历，改用向量化操作批量处理所有行的分段（可结合np.split和广播实现）
• 使用numba对核心函数进行JIT编译，加速循环和计算过程

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