如何不使用for循环填充二维二进制Numpy数组？

假设我有一个Numpy数组a，想要将其内部填充为全1，得到数组b：

原数组a：

print(a)
array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
       [0., 1., 0., 1., 1., 1., 1., 1., 0., 0.],
       [0., 1., 0., 0., 1., 0., 0., 0., 1., 0.],
       [0., 1., 0., 1., 0., 0., 0., 0., 1., 0.],
       [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
       [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
       [0., 0., 1., 0., 0., 0., 0., 1., 0., 0.],
       [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

目标数组b：

print(b)
array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 1., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
       [0., 1., 1., 1., 1., 1., 1., 1., 0., 0.],
       [0., 0., 1., 1., 1., 1., 1., 1., 0., 0.],
       [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

目前我用for循环实现的代码如下：

b = np.zeros(a.shape)
for i in range(a.shape[0]):
    occupied = np.where(a[i] == 1)[0]
    if len(occupied) > 0:
        for j in range(occupied[0], occupied[-1] + 1):
            b[i][j] = 1

请问是否有办法不使用for循环，仅通过Numpy完成此操作？

补充说明：

• 仅使用Numpy
• 需要填充的区域边界始终是连续的

可以通过Numpy的向量化操作完全替代循环，核心是定位每行中第一个和最后一个1的位置，再生成掩码填充中间区域：

实现步骤 #

1. 获取每行第一个1的索引
   用np.argmax找到每行第一个1的位置，同时处理全0行（将索引设为数组列数，避免后续误匹配）：

   first_ones = np.argmax(a == 1, axis=1)
   first_ones[~np.any(a == 1, axis=1)] = a.shape[1]
   

2. 获取每行最后一个1的索引
   翻转数组后用np.argmax找到最后一个1的位置，再转换为原数组的索引，同样处理全0行：

   last_ones = a.shape[1] - 1 - np.argmax(np.flip(a == 1, axis=1), axis=1)
   last_ones[~np.any(a == 1, axis=1)] = -1
   

3. 生成掩码并创建目标数组
   生成与原数组同形状的列索引矩阵，通过比较列索引是否在首尾1的区间内生成掩码，最后用np.where生成目标数组：

   col_indices = np.tile(np.arange(a.shape[1]), (a.shape[0], 1))
   mask = (col_indices >= first_ones[:, np.newaxis]) & (col_indices <= last_ones[:, np.newaxis])
   b = np.where(mask, 1.0, 0.0)
   

完整代码 #

import numpy as np

# 示例数组a
a = np.array([
    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
    [0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
    [0., 1., 0., 1., 1., 1., 1., 1., 0., 0.],
    [0., 1., 0., 0., 1., 0., 0., 0., 1., 0.],
    [0., 1., 0., 1., 0., 0., 0., 0., 1., 0.],
    [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
    [0., 1., 0., 1., 0., 0., 0., 1., 0., 0.],
    [0., 0., 1., 0., 0., 0., 0., 1., 0., 0.],
    [0., 0., 0., 1., 1., 1., 1., 0., 0., 0.],
    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]
])

# 获取每行第一个1的索引
first_ones = np.argmax(a == 1, axis=1)
first_ones[~np.any(a == 1, axis=1)] = a.shape[1]

# 获取每行最后一个1的索引
last_ones = a.shape[1] - 1 - np.argmax(np.flip(a == 1, axis=1), axis=1)
last_ones[~np.any(a == 1, axis=1)] = -1

# 生成列索引矩阵
col_indices = np.tile(np.arange(a.shape[1]), (a.shape[0], 1))

# 生成掩码
mask = (col_indices >= first_ones[:, np.newaxis]) & (col_indices <= last_ones[:, np.newaxis])

# 创建目标数组b
b = np.where(mask, 1.0, 0.0)

print(b)

这段代码完全依赖Numpy向量化操作，避免了循环，处理大型数组时效率会显著高于循环实现。

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