非二次复杂度实现Pandas数据反转的字典赋值方案

场景与模拟数据 #

处理包含300万+观测值的数据集，模拟数据及初始化代码如下：

import pandas as pd
import numpy as np
import math

# 初始化模拟数据
data = [[2, [4], None], [4, [9,18,6], None], [6, [], 9],[7, [2], None],[9, [4], 7],[14, [18,6], 3],[18, [7], 1]]
  
# 创建DataFrame
df = pd.DataFrame(data, columns=['docdb', 'cited_docdb','fronteer'])

# 处理空值与计算distance列
df.replace(' NaN', np.NaN)
df['distance'] = np.where((df['fronteer'] >0), 0, math.inf)

原实现的问题 #

原代码通过遍历唯一cited_docdb并反复查询DataFrame生成字典，属于二次复杂度（O(k*n)，k为唯一引用项数量），处理百万级数据时效率极低：

keys= [k for k in df.explode('cited_docdb')['cited_docdb'].unique()]
values=[df.explode('cited_docdb').loc[df.explode('cited_docdb')['cited_docdb']== j, 'docdb'].to_list() for j in keys]

inverting_db = dict(zip(keys, values))

优化方案（线性复杂度） #

方案一：利用Pandas分组操作（推荐） #

通过explode展开引用列后直接分组聚合，全程线性遍历，复杂度O(n)：

# 展开引用列并过滤空值
exploded_df = df.explode('cited_docdb').dropna(subset=['cited_docdb'])
# 按被引用项分组，收集对应的docdb列表并转字典
inverting_db = exploded_df.groupby('cited_docdb')['docdb'].apply(list).to_dict()

方案二：手动循环填充字典 #

直接遍历原DataFrame的每一行，逐个填充字典，复杂度O(n + m)（m为所有引用项的总数量）：

inverting_db = {}

for _, row in df.iterrows():
    current_docdb = row['docdb']
    cited_list = row['cited_docdb']
    
    # 跳过空引用列表
    if not cited_list:
        continue
    
    # 遍历每个被引用项，更新字典
    for cited in cited_list:
        if cited not in inverting_db:
            inverting_db[cited] = []
        inverting_db[cited].append(current_docdb)

效果说明 #

两种优化方案均避免了原方法中重复遍历数据集的开销，在300万+数据量下，执行效率会有数量级的提升。

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