python - 如何根据pandas中的关系列将id收集到列表中

Question

我有一个像这样的数据框:

  id near_relation
0  A        [B, D]
1  B     [A, H, N]
2  C        [I, R]
3  D        [A, E]
4  E        [D, M]
5  F        [J, K]
6  J        [F, P]
7  P        [J, S]

我想将 id 收集到基于相邻关系 near_relation 列的列表中。例如A和D有相邻关系，D和E有相邻关系，所以A 和 E 应收集在一个列表中。所以我期望的结果如下:

        collect
0  [A, B, D, E]
1           [C]
2     [F, P, J]

我可以用循环来完成此操作，但我的数据框很大，我想要一个更优雅的解决方案。

添加循环解决方案

def add(x,x_dict):
    tmp_list = x.copy()
    last_len = 0
    while len(tmp_list)!=last_len:
        current_len = len(tmp_list)
        for m in tmp_list[last_len:]:
            if m in x_dict.keys():
                for n in x_dict[m]:
                    if n not in tmp_list:
                        tmp_list.append(n)
        last_len = current_len
    return set(tmp_list) & set(x_dict.keys())

df_dict = df.set_index('id').to_dict()
df['collect'] = df['near_relation'].apply(lambda x: add(x,df_dict['near_relation']))

然后我根据上述循环的结果对它们进行分组。

有人可以帮助我吗？提前致谢。

Answer 1

这是一个 NetworkX 方法。首先获取所有 id-near_relation 元组，这些元组将用作构建图表的边:

l = [(i.id,nb) for _, i in df.iterrows() for nb in i.near_relation]
# [('A', 'B'), ('A', 'D'), ('B', 'A'), ('B', 'H'), ('B', 'N')...

根据上面的边列表构建一个图并计算其 connected_components 。最后仅保留那些也出现在 id 列中的组件:

import networkx as nx
G=nx.Graph()
G.add_edges_from(l)
cc = nx.connected_components(G)
[i & set(df.id) for i in cc]

输出

[{'A', 'B', 'D', 'E'}, {'C'}, {'F', 'J', 'P'}]