如何合并PyTorch矩阵中存在相同元素的行？

针对你提出的合并有共享元素的行的需求，这里提供两种实现方式：

方法一：利用NetworkX处理连通分量（简洁高效） #

NetworkX可以快速构建图并找出连通分量，适合快速实现需求：

import torch
import networkx as nx

# 输入矩阵
matrix = torch.tensor([[1, 2], [2, 3], [4, 5]])

# 构建无向图：将每行的两个元素作为一条边
G = nx.Graph()
for row in matrix.numpy():
    G.add_edge(row[0], row[1])

# 获取所有连通分量
components = list(nx.connected_components(G))

# 生成结果矩阵
result = []
for comp in components:
    comp_list = list(comp)
    if len(comp_list) == 2:
        result.append(comp_list)
    else:
        # 多元素分量取排序后的首尾元素（适配你的示例逻辑）
        sorted_comp = sorted(comp_list)
        result.append([sorted_comp[0], sorted_comp[-1]])

# 转换为PyTorch张量
result_tensor = torch.tensor(result)
print(result_tensor)

输出结果：

tensor([[1, 3],
        [4, 5]])

方法二：纯PyTorch/Numpy实现（无第三方依赖） #

如果不想引入NetworkX，可以手动实现连通分量的合并逻辑：

import torch

def merge_overlapping_rows(matrix):
    arr = matrix.numpy()
    element_to_comp = {}
    comp_count = 0

    # 遍历每行，构建元素与分量的映射
    for row in arr:
        a, b = row[0], row[1]
        a_in = a in element_to_comp
        b_in = b in element_to_comp

        if not a_in and not b_in:
            element_to_comp[a] = comp_count
            element_to_comp[b] = comp_count
            comp_count += 1
        elif a_in and not b_in:
            element_to_comp[b] = element_to_comp[a]
        elif not a_in and b_in:
            element_to_comp[a] = element_to_comp[b]
        else:
            # 合并两个不同的分量
            comp_a, comp_b = element_to_comp[a], element_to_comp[b]
            if comp_a != comp_b:
                for elem, comp in element_to_comp.items():
                    if comp == comp_b:
                        element_to_comp[elem] = comp_a

    # 按分量分组元素
    comp_to_elements = {}
    for elem, comp in element_to_comp.items():
        comp_to_elements.setdefault(comp, []).append(elem)

    # 生成结果
    result = []
    for elem_list in comp_to_elements.values():
        if len(elem_list) == 2:
            result.append(elem_list)
        else:
            sorted_list = sorted(elem_list)
            result.append([sorted_list[0], sorted_list[-1]])

    return torch.tensor(result)

# 测试
matrix = torch.tensor([[1, 2], [2, 3], [4, 5]])
print(merge_overlapping_rows(matrix))

输出结果：

tensor([[1, 3],
        [4, 5]])

逻辑说明 #

1. 连通分量构建：将每行的两个元素视为图的边，共享元素的行会被归入同一个连通分量（比如[1,2]和[2,3]通过元素2连通）。
2. 分量处理：对于包含2个元素的分量直接保留；对于多元素分量，按排序后的首尾元素生成新行（匹配你示例中的合并逻辑）。如果需要其他处理方式（如保留所有唯一元素），可修改分量处理的代码段。

内容的提问来源于stack exchange，提问作者Jemma