如何用numpy广播/向量化操作替换for循环实现逐行np.repeat功能

你基于单行np.repeat操作的延伸思路是完全正确的，不需要逐行写Python for循环，只要利用np.repeat支持传入一维重复次数序列的特性，配合切分操作就能一次性完成所有计算，性能远高于原循环实现：

# 完整可运行代码
# 1. 构造和rnd_ar展平后顺序一一对应的标签序列
# 展平后的rnd_ar顺序为[行0的O计数, 行0的T计数, 行1的O计数, 行1的T计数,...]
# 对应标签序列按同样顺序展开，和每个计数项一一匹配
labels = np.tile(['O', 'T'], (len(rnd_ar), 1)).ravel()
# 2. 一次性完成所有元素的重复操作，不需要逐行调用接口
all_repeated_seq = np.repeat(labels, rnd_ar.ravel())
# 3. 计算每一行的总长度，按累积切分点拆分出独立数组
row_lengths = rnd_ar.sum(axis=1)
split_idx = np.cumsum(row_lengths)
test = np.split(all_repeated_seq, split_idx)[:-1]

• 输出结果和原for循环生成的test完全一致：每个数组的内容、字符串类型、空数组格式都没有差异
• 针对你给出的样例rnd_ar，切分后的片段长度依次为6、3、2、7、1、2、1、0、3、1，和原循环的输出数组长度完全匹配
• 切分操作要注意：np.split会在最后一个切分点额外生成一个空数组，因此需要用[:-1]去掉最后这个冗余片段，保证列表长度和原rnd_ar的行数一致

如果可以接受逐行调用numpy接口的实现，用pandas的逐行处理代码会更直观，不需要手动计算切分点：

import pandas as pd
df = pd.DataFrame(rnd_ar, columns=['o_repeat', 't_repeat'])
test = df.apply(
    lambda row: np.repeat(['O', 'T'], [row.o_repeat, row.t_repeat]),
    axis=1
).tolist()

• 这种实现的逻辑和你写的for循环完全等价，只是把循环逻辑封装到了pandas的apply接口里
• 缺点是当数据量较大时，性能比纯NumPy向量化方案差很多，适合数据量小、追求代码可读性的场景

• 不要直接给np.repeat传入二维的rnd_ar，否则输出结果会变成二维数组，和原循环输出的一维数组列表格式不匹配
• 切分点必须用每一行的元素和做累积，不能直接用展平后的rnd_ar做切分，否则会把同一行的O、T片段错误拆分
• 注意保留字符串数组的类型，不要在拼接过程中把空数组变成其他类型，否则会和原循环的输出格式出现细微差异

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