优化Pandas循环效率：为8万问题生成含随机标题的新DataFrame

问题背景 #

现有一个Pandas DataFrame，包含约8万条唯一question_string（问题）、约8千条唯一titles（标题）。需求为：为每个问题随机选取200个不重复标题，其中1个为对应正确标题（标记is_similar=1），其余199个为错误标题（标记is_similar=0），最终生成新的DataFrame。原方案使用列表+for循环实现，处理速度极慢，需优化。

原始代码的性能问题 #

原代码存在以下低效点：

• 每次循环重复计算new_df.titles.unique().tolist()，8万次重复计算浪费大量资源。
• 使用Python列表的remove()和random.choices（未保证不重复）操作，列表本身的线性查找/修改效率极低。
• 累计1600万次full_list.append()操作，频繁触发列表内存扩容，耗时严重。

优化方案（基于Numpy+Pandas批量操作） #

利用Numpy的高效数组操作替代Python循环，减少内存开销和计算冗余：

import pandas as pd
import numpy as np

# 提前提取所有唯一标题、问题、索引数组
unique_titles = new_df['titles'].unique()
questions = new_df['question_string'].values
image_indexes = new_df['image_index'].values
correct_titles = new_df['titles'].values
n_samples = len(questions)
n_neg_per_question = 199  # 每个问题的负样本数量

# ---------------------- 生成正样本 ----------------------
pos_df = pd.DataFrame({
    'questions': questions,
    'titles': correct_titles,
    'image_index': image_indexes,
    'is_similar': np.ones(n_samples, dtype=np.int8)  # 用int8节省内存
})

# ---------------------- 批量生成负样本 ----------------------
# 初始化负样本数组，提前分配内存避免频繁扩容
neg_questions = np.empty(n_samples * n_neg_per_question, dtype=object)
neg_titles = np.empty(n_samples * n_neg_per_question, dtype=object)
neg_image_indexes = np.empty(n_samples * n_neg_per_question, dtype=image_indexes.dtype)
neg_is_similar = np.zeros(n_samples * n_neg_per_question, dtype=np.int8)

for i in range(n_samples):
    # 排除当前问题对应的正确标题
    valid_titles_mask = unique_titles != correct_titles[i]
    available_titles = unique_titles[valid_titles_mask]
    
    # 随机选择199个不重复的错误标题
    selected_neg_titles = np.random.choice(available_titles, size=n_neg_per_question, replace=False)
    
    # 填充数组（批量赋值，比逐个append高效）
    start_idx = i * n_neg_per_question
    end_idx = start_idx + n_neg_per_question
    neg_questions[start_idx:end_idx] = questions[i]
    neg_titles[start_idx:end_idx] = selected_neg_titles
    neg_image_indexes[start_idx:end_idx] = image_indexes[i]

# 构造负样本DataFrame
neg_df = pd.DataFrame({
    'questions': neg_questions,
    'titles': neg_titles,
    'image_index': neg_image_indexes,
    'is_similar': neg_is_similar
})

# ---------------------- 合并正负样本 ----------------------
full_list_df = pd.concat([pos_df, neg_df], ignore_index=True)

优化说明 #

1. 提前计算唯一标题：仅计算一次unique_titles，避免8万次重复计算。
2. Numpy数组批量操作：提前分配数组内存，批量填充数据，避免Python列表的频繁扩容开销，Numpy的随机选择np.random.choice是C级实现，远快于Python原生的random模块。
3. 内存优化：使用np.int8存储is_similar字段，减少内存占用（从int64降到int8，节省8倍内存）。
4. 保证不重复：np.random.choice设置replace=False，确保每个问题的199个错误标题无重复。

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