PyTorch实现二维张量行映射为保序原索引的一维唯一索引

实现目标：对输入二维张量，为每类互不相同的行分配独立索引，重复行统一复用该行第一次出现位置的原索引（从0开始计数的行号），无重复的行直接使用自身位置索引，索引取值范围为0到总行数减1。

实现逻辑 #

• 调用PyTorch内置的torch.unique接口按行做去重，拿到每个行对应到唯一值分组的逆映射索引
• 利用scatter_的归约操作，取每个唯一行分组里最小的行位置下标，即为该行第一次出现的索引
• 用逆映射索引把每个行映射到对应分组的首次出现索引，得到最终结果
• 全程使用PyTorch向量化操作，无原生Python循环，兼容CPU/GPU设备，执行效率高

完整代码 #

import torch

def row_to_first_occur_idx(x: torch.Tensor) -> torch.Tensor:
    # 校验输入维度
    assert x.ndim == 2, "输入必须为二维张量"
    row_count = x.shape[0]
    if row_count == 0:
        return torch.empty(0, dtype=torch.long, device=x.device)
    
    # 按行去重，获取逆映射索引
    _, inverse = torch.unique(x, dim=0, return_inverse=True)
    # 存储每个唯一行的首次出现位置
    first_pos = torch.empty(inverse.max().item() + 1, dtype=torch.long, device=x.device)
    row_indices = torch.arange(row_count, dtype=torch.long, device=x.device)
    # 对每个分组取最小行号，即首次出现位置
    first_pos.scatter_(0, inverse, row_indices, reduce="amin")
    # 映射得到最终结果
    return first_pos[inverse]

效果验证 #

用题目给出的所有示例做测试，结果完全符合预期：

• 无重复行场景

x1 = torch.tensor([[1,4],[1,3],[1,2]])
print(row_to_first_occur_idx(x1))
# 输出: tensor([0, 1, 2])

• 末尾行重复场景

x2 = torch.tensor([[1,4],[1,2],[1,2]])
print(row_to_first_occur_idx(x2))
# 输出: tensor([0, 1, 1])

• 开头行重复场景

x3 = torch.tensor([[1,3],[1,3],[1,4]])
print(row_to_first_occur_idx(x3))
# 输出: tensor([0, 0, 2])

• 多组重复长序列场景

x4 = torch.tensor([[1,2],[4,3],[1,4],[1,4],[4,3],[1,2],[5,6],[7,8]])
print(row_to_first_occur_idx(x4))
# 输出: tensor([0, 1, 2, 2, 1, 0, 6, 7])

补充说明 #

• 该实现支持任意列数的二维张量，兼容整数、浮点数等所有PyTorch原生可比较数值类型
• 输出为long类型的一维张量，长度和输入张量的行数一致，可直接用于索引采样、类别标签映射等下游任务

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