Python 如何合并具有相同元素的 torch 矩阵的行

发布于03月20日

例如，N * 2个输入matrix = torch.tensor([[1, 2], [2, 3],[4, 5]])，由于第一行和第二行共享相同的元素2，所以期望的输出矩阵是torch.tensor([[1, 3],[4, 5]]).

如何做到这一点？先谢谢你.

推荐答案

您可以标识0中的值与1中的前一行不同的行，并丢弃其他行:

mask = matrix[1:,0]!=matrix[:-1,1]
# tensor([False,  True])

true = torch.tensor([True])

out = torch.column_stack([matrix[torch.cat((true, mask)), 0],
                          matrix[torch.cat((mask, true)), 1],
                         ])

变式:

mask = torch.cat((torch.tensor([True]),
                  matrix[1:,0]!=matrix[:-1,1]))
# tensor([ True, False,  True])

out = torch.column_stack([matrix[mask, 0],
                          matrix[mask.roll(-1), 1],
                         ])

输出:

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

更复杂的例子:

# input             # values to keep (X)
tensor([[1, 2],     #    X   -
        [2, 3],     #    -   -
        [3, 4],     #    -   X
        [0, 0],     #    X   X
        [5, 6],     #    X   -
        [6, 9]])    #    -   X

# output
tensor([[1, 4],
        [0, 0],
        [5, 9]])