Как я могу собрать идентификатор в список в соответствии с колонкой отношений в пандах

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