DataFrame内存高效循环优化：重量匹配场景内存溢出解决

现有如下DataFrame：

原代码通过循环匹配重量偏差50以内的产品，同时包含weight_in_g=0的项（占数据集35%），代码如下：

list1=[]

for row in df[['product_id', 'weight_in_g']].itertuples():
    high = row[1] + 50
    low = row[1] - 50 

    id = df['product_id'].loc[((df['weight_in_g'] >= low) & (df['weight_in_g'] <= high)) | (df['weight_in_g'] == 0)]
    list1.append(id)
    
df['weight_matches'] = list1
del list1

运行后输出符合预期，但数据集超过2万行时，13GB内存的笔记本会出现内存耗尽问题，需要更内存高效的处理方案。

方案1：排序+滑动窗口（大规模数据集首选） #

原循环的核心问题是每次迭代全表扫描，时间复杂度O(n²)，且重复生成大量Series对象占用内存。先按重量排序，再用滑动窗口快速定位匹配区间，能把时间和内存复杂度降到O(n log n)。

代码示例：

import pandas as pd
import numpy as np

# 1. 一次性提取所有weight_in_g=0的产品ID，避免重复过滤
zero_ids = df[df['weight_in_g'] == 0]['product_id'].tolist()

# 2. 分离非0数据并按重量排序
non_zero_df = df[df['weight_in_g'] != 0].sort_values('weight_in_g').reset_index(drop=True)
weights = non_zero_df['weight_in_g'].values
product_ids = non_zero_df['product_id'].values

# 3. 用numpy.searchsorted快速定位每个重量的上下区间索引
left_indices = np.searchsorted(weights, weights - 50, side='left')
right_indices = np.searchsorted(weights, weights + 50, side='right')

# 4. 生成每个产品的匹配ID（合并0重量项）
matches = []
for l, r in zip(left_indices, right_indices):
    interval_ids = product_ids[l:r].tolist()
    full_ids = interval_ids + zero_ids
    matches.append(full_ids)

non_zero_df['weight_matches'] = matches

# 5. 单独处理weight_in_g=0的行：匹配所有重量在-50~50之间的产品（含自身）
zero_df = df[df['weight_in_g'] == 0].copy()
zero_match_ids = df[(df['weight_in_g'] >= -50) & (df['weight_in_g'] <= 50)]['product_id'].tolist()
zero_df['weight_matches'] = [zero_match_ids] * len(zero_df)

# 6. 合并结果并还原原排序
final_df = pd.concat([non_zero_df, zero_df]).sort_values('product_id').reset_index(drop=True)

方案2：向量化操作（中等规模数据集适用） #

通过广播生成重量差值矩阵，减少循环中的重复计算，但需注意：如果数据集超过10万行，布尔矩阵会占用较多内存，此时优先用方案1。

代码示例：

import pandas as pd
import numpy as np

zero_ids = df[df['weight_in_g'] == 0]['product_id'].tolist()
non_zero_df = df[df['weight_in_g'] != 0].copy()

# 生成重量差值矩阵，标记差值≤50的项
weights = non_zero_df['weight_in_g'].values
diff_matrix = np.abs(weights[:, None] - weights) <= 50

# 提取每行符合条件的ID并合并0重量项
non_zero_df['weight_matches'] = [
    non_zero_df['product_id'][mask].tolist() + zero_ids
    for mask in diff_matrix
]

# 处理0重量行
zero_df = df[df['weight_in_g'] == 0].copy()
zero_match_mask = np.abs(df['weight_in_g'].values - 0) <= 50
zero_df['weight_matches'] = [df['product_id'][zero_match_mask].tolist()] * len(zero_df)

final_df = pd.concat([non_zero_df, zero_df]).sort_values('product_id').reset_index(drop=True)

核心优化点 #

• 避免重复扫描：排序+滑动窗口将每次O(n)的全表扫描改为O(log n)的索引查找
• 减少内存冗余：用列表存储匹配ID，替代原代码中重复生成的Series对象
• 批量处理0项：一次性提取所有0重量ID，避免循环中重复过滤

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