高效实现DataFrame值与另一DataFrame多列匹配生成嵌套列表

需求说明 #

针对df1中每个NID列的值，找出其在df2的N1、N2、N3、N4列中出现时对应的EID值，生成嵌套列表。

示例数据 #

# df1
   NID
0    1
1    2
2    3
3    4
4    5

# df2
   EID  N1  N2  N3  N4
0    1   1   2  13  12
1    2   2   3  14  13
2    3   3   4  15  14
3    4   4   5  16  15
4    5   5   6  17  16
5    6   6   7  18  17
6    7   7   8  19  18
7    8   8   9  20  19
8    9   9  10  21  20
9   10  10  11  22  21

预期结果 #

sol = [[1], [1, 2], [2, 3], [3, 4], [4, 5]]

现有低效实现（双重循环） #

当前采用双重循环实现，但数据量达数千条时效率极低，代码如下：

import pandas as pd
df1 = pd.DataFrame({'NID':[1,2,3,4,5]})
df2 = pd.DataFrame({'EID':[1,2,3,4,5,6,7,8,9,10],
                  'N1':[1,2,3,4,5,6,7,8,9,10],
                  'N2':[2,3,4,5,6,7,8,9,10,11],
                  'N3':[13,14,15,16,17,18,19,20,21,22],
                  'N4':[12,13,14,15,16,17,18,19,20,21]})

sol= []
for idx,node in df1.iterrows():
    x = []
    for idx2,elem in df2.iterrows():
        if node['NID'] == elem['N1']:
            x.append(elem['EID'])
        if node['NID'] == elem['N2']:
            x.append(elem['EID'])
        if node['NID'] == elem['N3']:
            x.append(elem['EID'])
        if node['NID'] == elem['N4']:
            x.append(elem['EID'])
    sol.append(x)

print(sol)

高效无循环实现方案 #

利用Pandas的melt函数将df2转换为长格式，再通过分组和映射完成匹配，避免循环，大幅提升效率：

实现代码 #

import pandas as pd

# 转换df2为长格式：将N1-N4列合并为一列，保留对应的EID
melted_df = df2.melt(id_vars='EID', value_vars=['N1', 'N2', 'N3', 'N4'], var_name='N_col', value_name='N_value')

# 按N_value分组，收集对应的EID列表，生成字典映射
nid_to_eids = melted_df.groupby('N_value')['EID'].apply(list).to_dict()

# 针对df1的每个NID，从字典中获取对应的EID列表，不存在则返回空列表
sol = [nid_to_eids.get(nid, []) for nid in df1['NID']]

print(sol)

代码说明 #

1. melt转换长格式：将df2的宽格式（多N列）转为长格式，每一行对应一个EID和一个N值，方便后续分组操作。
2. 分组生成映射字典：按N_value分组，将每组的EID收集为列表，转为字典后可以O(1)时间查询每个NID对应的EID列表。
3. 映射生成结果：遍历df1的NID，通过字典快速获取对应EID列表，生成最终的嵌套列表。

这种方法的时间复杂度远低于双重循环，在大数据量场景下性能提升显著。

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