如何根据熊猫中的关系列将ID收集到列表中

时间:2019-02-25 13:36:00

标签: python pandas dataframe networkx

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

  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收集到列表中。例如,AD具有相邻关系,DE具有相邻关系,因此AE应该收集在一个清单。所以我期望的结果如下:

        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']))

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

有人可以帮助我吗?预先感谢。

1 个答案:

答案 0 :(得分:2)

这是一种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'}]