基于值偏移规则创建布尔数组的Numpy高效实现方法

给定整数d和布尔数组M，示例如下：

import numpy as np
M = np.array([[ True, False, False,  True],
              [False,  True, False, False],
              [False, False, False, False],
              [False, False, False,  True]])

需求：生成新数组，保留原数组的True，且每个True上方最多d行的同一列位置也设为True。已用双重循环实现，现寻求Numpy原生高效方案。

补充示例：

• 当d=1时，预期结果：

np.array([[ True, False,  True,  True],
          [ True,  True, False, False],
          [False, False, False, False],
          [False, False,  True,  True]])

• 当d=2时，预期结果：

np.array([[ True,  True,  True,  True],
          [ True,  True, False, False],
          [False, False, False, False],
          [False,  True,  True,  True]])

方案1：基于广播的半向量化实现 #

该方案通过遍历每个True的位置，利用Numpy切片批量设置目标行，比双重循环效率更高：

def expand_true(M, d):
    result = M.copy()
    true_rows, true_cols = np.where(M)
    for r, c in zip(true_rows, true_cols):
        start = max(0, r - d)
        if start < r:
            result[start:r, c] = True
    return result

方案2：完全向量化实现（无显式循环） #

利用矩阵乘法生成扩展掩码，结合布尔或运算得到最终结果：

def expand_true(M, d):
    n_rows = M.shape[0]
    M_int = M.astype(np.int8)
    
    # 构建标记矩阵：tri[i,j] = 1 表示行i属于行j的上方d行范围内
    tri = np.zeros((n_rows, n_rows), dtype=np.int8)
    for i in range(n_rows):
        start = max(0, i - d)
        tri[start:i, i] = 1
    
    # 计算扩展的True区域，合并原数组的True
    expanded = (tri @ M_int) > 0
    return M | expanded

方案3：利用一维卷积（需scipy） #

通过列方向的滑动窗口最大值实现扩展，代码更简洁：

from scipy.ndimage import convolve1d

def expand_true(M, d):
    # 卷积核长度为d+1，覆盖当前行及上方d行
    kernel = np.ones(d + 1, dtype=int)
    # 对每列做卷积，结果大于0即表示存在True在范围内
    expanded = convolve1d(M.astype(int), kernel, axis=0, mode='constant', cval=0) > 0
    return expanded

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