如何提升Pandas DataFrame中随机列选择与赋值的效率？

你的核心问题在于使用了iterrows()逐行迭代，这在处理大型DataFrame时会因为Python级别的循环导致性能瓶颈，同时代码里还有冗余的列复制操作。下面是针对性的优化方案：

优化思路 #

1. 移除冗余操作：new_df = df.copy()已经完整复制了原DataFrame的所有列，不需要再循环赋值非相关列。
2. 向量化计算概率：用pandas的向量化方法代替逐行计算概率，避免循环开销。
3. 批量生成随机选择结果：利用numpy的向量化随机函数，一次性完成所有行的列选择，彻底摆脱逐行循环。

优化后的完整代码 #

import pandas as pd
import numpy as np

# 示例DataFrame
data = {
    'dog': [1, 2, 3, 4],
    'cat': [5, 6, 7, 8],
    'parrot': [9, 10, 11, 12],
    'owner': ['fred', 'bob', 'jim', 'jannet']
}
df = pd.DataFrame(data)

relevant_col_list = ['dog', 'cat', 'parrot']

# 直接复制原DataFrame，无需后续循环复制列
new_df = df.copy()

# 1. 向量化计算每行的概率矩阵
row_sums = df[relevant_col_list].sum(axis=1)
prob_matrix = df[relevant_col_list].div(row_sums, axis=0)

# 2. 批量生成随机选择结果
rng = np.random.default_rng()  # 使用更现代的随机数生成器
rand_vals = rng.random(len(df))
# 计算累积概率，找到每行第一个大于随机值的列
cumulative_probs = prob_matrix.cumsum(axis=1)
chosen_columns = cumulative_probs.ge(rand_vals[:, np.newaxis]).idxmax(axis=1)

# 赋值新列
new_df['iteration_1'] = chosen_columns

print(new_df)

性能提升说明 #

• 向量化操作：pandas和numpy的底层是C实现的向量化计算，速度比Python循环快几个数量级，尤其当数据集行数超过1万时，差异会非常明显。
• 无冗余操作：删除了原代码中重复复制列的无效逻辑，减少不必要的内存和CPU开销。
• 批量随机选择：通过累积概率与随机值的比较，一次性完成所有行的列选择，避免了逐行调用np.random.choice的开销。

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