将数据帧中的值与另一个数据帧的多列进行比较，以获取条目以有效方式匹配的列表列表

Question

我有两个 pandas 数据帧，我想找到第二个数据帧中出现特定值的所有条目。

举个例子：

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

现在，我基本上想要的是具有值 EID（来自 df2）的列表列表，其中值 NID（来自 df1）出现在任何列 N1、N2、N3、N4 中：

解决方案是：

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

所需的解决方案解释：

该解决方案有 5 个条目 (len(sol = 5))，因为我在 df1 中有 5 个条目。sol 中的第一个条目是 1，因为对于 df2 中的 EID=1，值 NID = 1 仅出现在 N1,N2,N3,N4 列中。sol 中的第二个条目引用值 NID=2（df1）并且长度为 2，因为可以在 N1 列（对于 EID=2）和 N2 列（对于 EID=1）中找到 NID=2。因此，解中的第二个条目是[1,2]，依此类推。

到目前为止，我尝试的是循环 df1 中的每个元素，然后循环 df2 中的每个元素以查看 NID 是否位于 N1、N2、N3、N4 列中的任何一列中。该解决方案有效，但对于巨大的数据帧（每个 df 最多可以有数千个条目），该解决方案变得极其耗时。因此我一直在寻找一种更有效的解决方案。

我的代码已实现：

输入数据：

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)

如果有人有一个不需要循环的解决方案，我会很高兴。也许使用 numpy 函数或类似 cKDTrees 的东西，但不幸的是我不知道如何以更快的方式解决这个问题。

先感谢您！

Answer 1

您可以使用 重塑形状melt、使用 过滤loc以及groupby.agg作为列表。然后reindex转换tolist：

out = (df2
    .melt('EID')   # reshape to long form
    # filter the values that are in df1['NID']
    .loc[lambda d: d['value'].isin(df1['NID'])]
    # aggregate as list
    .groupby('value')['EID'].agg(list)
    # ensure all original NID are present in order
    # and convert to list
    .reindex(df1['NID']).tolist()
)

替代方案stack：

df3 = df2.set_index('EID')

out = (df3
    .where(df3.isin(df1['NID'].tolist())).stack()
    .reset_index(name='group')
    .groupby('group')['EID'].agg(list)
    .reindex(df1['NID']).tolist()
)

输出：

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