如何在高维NumPy布尔数组中沿轴传播True值？

需要将布尔数组中的True值向后传播指定数量的位置（示例中为2个位置），一维实现已存在但无法直接复用至高维场景，以下提供两种通用的向量化解决方案：

方案一：基于多维索引的向量化实现 #

该方法通过生成所有有效传播索引直接赋值，性能最优且适用于任意维度：

import numpy as np

propagation = 2
AA = np.array([
    [False, True, False, False, False, True, False, False, False, False, False, True, False],
    [False, True, False, False, False, True, False, False, False, False, False, True, False]
])

# 获取所有True值的多维坐标
true_coords = np.argwhere(AA)
# 生成传播方向的偏移量（0表示原位置，1到propagation表示后续位置）
offsets = np.arange(propagation + 1)

# 构造偏移矩阵：仅在最后一维（列）添加偏移，其他维度偏移为0
offset_matrix = np.zeros((propagation + 1, AA.ndim), dtype=int)
offset_matrix[:, -1] = offsets

# 广播偏移到每个True坐标，生成所有需要设为True的坐标
expanded_coords = true_coords[:, None, :] + offset_matrix
# 过滤掉超出数组边界的无效坐标
valid_mask = np.all(expanded_coords < np.array(AA.shape), axis=-1)
valid_coords = expanded_coords[valid_mask]

# 初始化结果数组并赋值
BB = np.zeros_like(AA)
BB[tuple(valid_coords.T)] = True

# 输出结果
print("原数组 AA:")
print(AA)
print("\n传播后的数组 BB:")
print(BB)

方案说明 #

• 支持任意维度数组：只需修改offset_matrix中偏移对应的轴即可（比如要在第0维传播，就把offset_matrix[:, 0] = offsets）。
• 无显式循环：全程使用NumPy广播机制，性能接近原生C实现。
• 自动边界处理：通过np.all(expanded_coords < AA.shape)过滤越界索引，避免报错。

方案二：基于累积最大值的简洁实现 #

如果传播步数不大，可使用反向累积最大值的方式实现，代码更简洁：

def propagate_boolean(arr, steps):
    # 反转目标维度，将"向后传播"转为"向前累积"
    reversed_arr = arr[..., ::-1]
    # 逐步累积最大值，实现传播效果
    for _ in range(steps):
        # 向右滚动一位，与原数组取最大值
        rolled = np.roll(reversed_arr, 1, axis=-1)
        # 滚动后首位会被填充原最后一位，这里重置为原数组的最后一位（避免错误传播）
        rolled[..., 0] = reversed_arr[..., 0]
        reversed_arr = np.maximum(reversed_arr, rolled)
    # 反转回原顺序
    return reversed_arr[..., ::-1]

BB = propagate_boolean(AA, propagation)
print("\n方案二输出结果:")
print(BB)

方案说明 #

• 代码简洁易读，适合小传播步数场景。
• 若传播步数较大，循环会带来性能损耗，此时方案一更优。

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