Python Numpy按首列权重重复行实现数组反扁平化方法

核心实现代码 #

不需要手写循环，直接使用NumPy内置的repeat函数处理整个二维数组即可，运行效率远高于手写循环：

import numpy as np

# 提取首列权重值，转为整数类型作为每行的重复次数
# 如果权重是带小数的浮点数，可替换为np.round(chain[:, 0]).astype(int)做四舍五入取整
repeat_counts = chain[:, 0].astype(int)
# 沿行维度(axis=0)按对应次数重复每一行
ACT_chain = np.repeat(chain, repeats=repeat_counts, axis=0)

效果验证 #

用你提供的示例数组测试：

# 构造示例输入
chain = np.array([
    [1, 5, 3],
    [2, 2, 1],
    [3, 0, 1]
])
repeat_counts = chain[:, 0].astype(int)
result = np.repeat(chain, repeats=repeat_counts, axis=0)
print(result)

输出完全匹配预期：

[[1 5 3]
 [2 2 1]
 [2 2 1]
 [3 0 1]
 [3 0 1]
 [3 0 1]]

之前写法的错误点 #

• 对一维单行数组调用np.repeat：一维数组不存在行维度，axis=0会沿元素维度逐值重复，因此出现每个浮点数连续重复N次的错误输出。
• 手写循环逻辑错误：无论重复次数多少都会先追加1次原行，重复次数>1时又追加np.tile生成的N行二维数组，既会导致行计数超出预期，还会造成列表内混合一维、二维数组的嵌套问题，无法生成规则的二维数组。

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