如何高效生成每行含3个长度100连续掩码块的随机mask矩阵

原代码效率低的核心原因是使用了三层Python层循环，没有利用Numpy的向量化运算能力，下面给出两种优化实现：

基础高效实现（允许掩码块重叠） #

完全基于Numpy向量化操作，无Python层循环，耗时为毫秒级：

import numpy as np

n_rows, n_cols = 10000, 1000
block_length = 100
blocks_per_row = 3

# 初始化全1掩码，用布尔类型更省内存、运算更快
mask = np.ones((n_rows, n_cols), dtype=np.bool_)
# 生成每行3个块的起始索引
start_indices = np.random.choice(n_cols - block_length + 1, size=(n_rows, blocks_per_row))
# 生成块内偏移量
offset = np.arange(block_length)
# 广播得到所有需要置0的列坐标，形状为 [10000, 3, 100]
col_pos = start_indices[..., None] + offset
# 构造对应形状的行坐标
row_pos = np.arange(n_rows)[:, None, None]
# 批量赋值为0
mask[row_pos, col_pos] = 0

进阶实现（保证每行3个掩码块互不重叠） #

采用隔板法分配块间距，从根源上避免块重叠，性能和基础版几乎一致：

import numpy as np

n_rows, n_cols = 10000, 1000
block_length = 100
blocks_per_row = 3
# 计算所有块占完后剩余的可分配间隔总长度
total_free_length = n_cols - blocks_per_row * block_length

mask = np.ones((n_rows, n_cols), dtype=np.bool_)
# 生成每行4个间隔（块前、块1与块2之间、块2与块3之间、块后），总和等于总空闲长度
gaps = np.random.multinomial(total_free_length, np.ones(4)/4, size=n_rows)
# 累加间隔和块长度，得到不重叠的块起始索引
start_indices = np.cumsum(gaps[:, :-1] + [0, block_length, block_length], axis=1)
# 后续赋值逻辑和基础版一致
offset = np.arange(block_length)
col_pos = start_indices[..., None] + offset
row_pos = np.arange(n_rows)[:, None, None]
mask[row_pos, col_pos] = 0

你可以验证任意一行的块起始位置，都会满足start_indices[i][1] >= start_indices[i][0] + 100、start_indices[i][2] >= start_indices[i][1] + 100，天然保证无重叠。

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