将概率转二进制特征：百万行DataFrame向量化优化需求

当然可以！用行循环处理百万级数据确实会慢到让人崩溃，我们可以借助NumPy和Pandas的向量化能力彻底解决这个问题，把速度提升几个数量级。核心思路是用底层数组运算替代Python层面的循环，充分利用CPU的向量计算优势。

优化后的完整代码 #

import pandas as pd
import numpy as np

np.random.seed(42)
df = pd.DataFrame({"T1" : [0.8,0.5,0.01],"T2":[0.1,0.2,0.89],"T3":[0.1,0.3,0.1]})

# 1. 计算每行的累积概率分布
cumulative_probs = df.values.cumsum(axis=1)
# 2. 为每行生成一个0-1之间的随机数（广播为列向量方便对比）
random_values = np.random.rand(df.shape[0], 1)
# 3. 找到每个随机数落在哪个概率区间，得到选中的列索引
selected_col_indices = (cumulative_probs >= random_values).argmax(axis=1)
# 4. 创建全0的二进制矩阵，将选中的位置设为1
binary_matrix = np.zeros_like(df.values)
binary_matrix[np.arange(df.shape[0]), selected_col_indices] = 1
# 5. 转换回DataFrame，保持原索引和列名
df_binary = pd.DataFrame(binary_matrix, columns=df.columns, index=df.index)

print(df_binary)

代码原理详解 #

1. 累积概率计算：cumsum(axis=1) 会对每行的概率累加，比如第一行的累积结果是 [0.8, 0.9, 1.0]，这样我们可以通过随机数落在哪个区间来判断选中的类别。
2. 随机数生成：np.random.rand 一次性生成所有行的随机数，避免逐行生成的开销。
3. 快速定位选中列：通过广播对比累积概率和随机数，argmax(axis=1) 会找到每行第一个满足条件的列索引（也就是随机数对应的类别），这一步是完全向量化的。
4. 构建二进制矩阵：利用NumPy的索引赋值，一次性将所有选中的位置设为1，其他保持0，比逐行逐列赋值高效得多。

效率对比 #

• 原循环方案：时间复杂度是 O(n*m)（n为行数，m为列数），且每一步都是Python层面的循环，百万行数据可能需要几分钟甚至更久。
• 向量化方案：同样是 O(n*m)，但所有运算都在C语言层面执行，百万行数据通常只需要几秒就能完成，速度提升至少100倍以上。

可选替代方案（稍逊但更直观） #

如果你觉得上面的步骤有点绕，也可以用np.random.multinomial实现，虽然比纯向量化稍慢，但也远快于循环：

# 生成每行的多项分布采样结果（每行只有一个1，其余0）
samples = np.array([np.random.multinomial(1, row) for row in df.values])
df_binary = pd.DataFrame(samples, columns=df.columns, index=df.index)

不过这个方法本质还是列表推导式的隐式循环，性能不如前面的纯向量化方案。

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