基于Numpy的二维0-1数组行列方向累计求和问题（无循环实现）

需求说明 #

给定由0和1组成的二维Numpy数组，需要为每个元素计算从自身位置开始，沿指定方向（行/列）到第一个0为止的所有元素之和，分别生成按列求和的数组col和按行求和的数组row。由于实际应用于大规模二进制图像，要求完全无循环实现以保证效率。

示例 #

输入数组 #

import numpy as np

inp = np.array([[0, 1, 1, 1, 0, 1, 1, 0, 1],
                [1, 1, 0, 0, 0, 1, 0, 1, 0],
                [1, 1, 1, 0, 1, 1, 0, 1, 1]])

按列求和输出（col） #

array([[0, 3, 2, 1, 0, 2, 1, 0, 1],
       [2, 1, 0, 0, 0, 1, 0, 1, 0],
       [3, 2, 1, 0, 2, 1, 0, 2, 1]])

按行求和输出（row） #

array([[0, 3, 1, 1, 0, 3, 1, 0, 1],
       [2, 2, 0, 0, 0, 2, 0, 2, 0],
       [1, 1, 1, 0, 1, 1, 0, 1, 1]])

无循环实现方案 #

核心思路是利用Numpy的向量运算和累积操作，定位每个元素到方向上最近的0的位置，再计算区间和。

1. 按行求和实现（mysum_rows） #

def mysum_rows(arr):
    # 复制数组并在每行末尾添加0，处理行尾无0的边界情况
    padded = np.hstack([arr, np.zeros((arr.shape[0], 1), dtype=arr.dtype)])
    # 提取每行中所有0的列位置
    zero_pos = np.where(padded == 0)[1].reshape(arr.shape[0], -1)
    
    # 生成行列索引矩阵，用于批量定位元素
    row_indices = np.arange(arr.shape[0])[:, None]
    col_indices = np.arange(arr.shape[1])[None, :]
    # 用二分查找定位每个元素右侧第一个0的位置
    first_zero_idx = np.array([np.searchsorted(zp, ci) for zp, ci in zip(zero_pos, col_indices)])
    first_zero_pos = zero_pos[row_indices, first_zero_idx]
    
    # 计算每行的累积和，通过累积和差值快速得到区间和
    row_cumsum = np.cumsum(arr, axis=1)
    # 处理当前元素为0的情况，直接赋值0
    result = np.where(arr == 0, 0, row_cumsum[:, first_zero_pos-1] - row_cumsum[:, col_indices-1] + arr[:, col_indices])
    # 单独处理行首元素的边界问题
    result[:, 0] = np.where(arr[:,0]==0, 0, row_cumsum[:, first_zero_pos[:,0]-1] if first_zero_pos[:,0]>0 else arr[:,0])
    
    return result

2. 按列求和实现（mysum_cols） #

将列方向的问题转置为行方向问题，复用行求和逻辑后再转置回原形状：

def mysum_cols(arr):
    return mysum_rows(arr.T).T

方案说明 #

• 用np.cumsum计算累积和，通过累积和的差值快速获取区间元素和，避免逐元素遍历。
• 借助np.searchsorted实现无循环的最近0位置查找，保证运算效率。
• 转置操作复用行求和逻辑，减少代码冗余，同时保持向量运算的高效性。

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