python - 有没有更优雅/优化的方法可以制作这种连接算法？

Question

我做了一个算法，从原子列表(给定它们彼此之间的距离)中计算出分子中连接了哪些原子。在物理环境之外考虑这个问题，它只是一个封闭的网络问题，节点是原子，边缘是连接原子的原子键。我有一个节点列表和一个连接节点的边列表，我需要找出每个独特分子的列表。我已经在下面的代码中完成了这个，但是它有点慢而且很丑陋。有没有办法优化这个算法？

这是我的代码，带有相关信息供您尝试(我将提供另一个原子列表来尝试调用 pairs_1 和 chosen_atom_1，只需将 pairs_1 更改为 pairs 并将 chosen_atom_1 更改为 chosen_atom 即可正常工作)

pairs = [[0, 1],
 [0, 2],
 [3, 4],
 [3, 5],
 [6, 7],
 [6, 8],
 [9, 10],
 [9, 11],
 [12, 13],
 [12, 14],
 [15, 16],
 [15, 17],
 [18, 19],
 [18, 20],
 [21, 22],
 [21, 23],
 [24, 25],
 [24, 26],
 [27, 28],
 [27, 29],
 [30, 31],
 [30, 32],
 [33, 34],
 [33, 35],
 [36, 37],
 [36, 38],
 [39, 40],
 [39, 41],
 [42, 43],
 [42, 44],
 [45, 46],
 [45, 47]]

chosen_atom = [np.random.rand() for i in range(48)]

pairs_1 = [[0, 6],
 [1, 7],
 [2, 8],
 [3, 9],
 [4, 10],
 [5, 6],
 [5, 10],
 [6, 7],
 [7, 8],
 [8, 9],
 [9, 10]]

chosen_atom_1 = [np.random.rand() for i in range(11)]

# use list of lists to define unique molecules
molecule_list = []

for i in pairs:
    temp_array = []
    for ii in pairs:
        temp_pair = [i[0], i[1]]

        if temp_pair[0] == ii[0]:
            temp_array.append(ii[1])
            temp_array = set(temp_array)
            temp_array = list(temp_array)
        if temp_pair[1] == ii[1]:
            temp_array.append(ii[0])
            temp_array = set(temp_array)
            temp_array = list(temp_array)

        for iii in temp_array:
            for j in pairs:
                if iii == j[0]:
                    temp_array.append(j[1])
                    temp_array = set(temp_array)
                    temp_array = list(temp_array)
                if iii == j[1]:
                    temp_array.append(j[0])
                    temp_array = set(temp_array)
                    temp_array = list(temp_array)

        if len(temp_array) > len(chosen_atom):
            break
    molecule_list.append(temp_array)

molecule_list = [list(item) for item in set(tuple(row) for row in molecule_list)]

# the output of pairs should be
molecule_list = [[8, 6, 7],
 [27, 28, 29],
 [9, 10, 11],
 [0, 1, 2],
 [32, 30, 31],
 [18, 19, 20],
 [45, 46, 47],
 [33, 34, 35],
 [24, 25, 26],
 [42, 43, 44],
 [16, 17, 15],
 [12, 13, 14],
 [21, 22, 23],
 [3, 4, 5],
 [40, 41, 39],
 [36, 37, 38]]
# the output of pairs_1 should be:
molecule_list = [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]

所以上面我已经给出了我现在得到的输出，并且应该得到 - 任何使这段代码更好的想法将不胜感激。

Answer 1

正如我在评论中所说，您需要连通分量算法，您可以使用 networkx 包轻松解决它:

import networkx as nx

G = nx.from_edgelist(pairs)

print([i for i in nx.connected_components(G)])

# jupyter notebook
%matplotlib inline
nx.draw(G, with_labels=True)

输出:

[{0, 1, 2}, {3, 4, 5}, {8, 6, 7}, {9, 10, 11}, {12, 13, 14}, {16, 17, 15}, {18, 19, 20}, {21, 22, 23}, {24, 25, 26}, {27, 28, 29}, {32, 30, 31}, {33, 34, 35}, {36, 37, 38}, {40, 41, 39}, {42, 43, 44}, {45, 46, 47}]

这是我的 script从原子坐标构建和可视化分子图。