PyTorch张量中高效查找唯一行首次出现索引的优化方法

import torch
import numpy as np

# 示例张量
data = torch.rand(100, 5)
# 随机复制一些行，制造重复数据
data[np.random.choice(100, 50, replace=False)] = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])

# 查找唯一行及其相关信息
# u_data: 唯一行张量
# inverse_indices: 原始张量中每行对应的唯一行索引
# counts: 每个唯一行出现的次数
u_data, inverse_indices, counts = torch.unique(data, dim=0, return_inverse=True, return_counts=True)

# 查找每个唯一行首次出现的索引（低效循环方法）
unique_indices_loop = torch.zeros(len(u_data), dtype=torch.long)
for idx in range(len(u_data)):
    # 对于每个唯一行索引idx，在inverse_indices中找到其首次出现的位置
    unique_indices_loop[idx] = torch.where(inverse_indices == idx)[0][0]

print(f"通过循环找到的首次出现索引: {unique_indices_loop}")

import torch
import numpy as np

# 示例张量（与前文相同）
data = torch.rand(100, 5)
data[np.random.choice(100, 50, replace=False)] = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])

# 查找唯一行及其相关信息
u_data, inverse_indices, counts = torch.unique(data, dim=0, return_inverse=True, return_counts=True)

# --- 优化的PyTorch方法 ---

# 1. 初始化辅助张量 A
# A 的维度为 (原始行数, 唯一行数)
# 使用一个大于原始行数的占位符值（如 len(data) 或 1000）进行初始化
# 这里使用 len(data) 作为占位符，因为原始行索引最大为 len(data) - 1
placeholder_val = len(data) 
A = placeholder_val * torch.ones((len(data), len(u_data)), dtype=torch.long)

# 2. 填充辅助张量 A
# 对于原始张量中的每一行 i (由 torch.arange(len(data)) 生成)，
# 找到其对应的唯一行索引 inverse_indices[i]。
# 然后在 A[i, inverse_indices[i]] 的位置上写入原始行索引 i。
# 这里的赋值操作 A[torch.arange(len(data)), inverse_indices] = torch.arange(len(data))
# 实际上是将原始行索引本身作为值写入，而不是 inverse_indices。
# 修正：应写入原始行索引 i，而不是 inverse_indices[i]
A[torch.arange(len(data)), inverse_indices] = torch.arange(len(data))

# 3. 对 A 沿着列方向执行 argmin 操作
# 对于 A 的每一列 j（代表第 j 个唯一行），argmin(A[:, j]) 将返回该列中最小值的行索引。
# 由于我们只在 A[i, j] 处写入了 i，且其他地方是更大的占位符，
# 因此 argmin 会找到第一个将自身映射到唯一行 j 的原始行索引。
unique_indices_optimized = torch.argmin(A, dim=0)

print(f"通过优化方法找到的首次出现索引: {unique_indices_optimized}")

# 验证两种方法的结果是否一致
# 为了验证，我们还需要运行之前的循环方法
unique_indices_loop = torch.zeros(len(u_data), dtype=torch.long)
for idx in range(len(u_data)):
    unique_indices_loop[idx] = torch.where(inverse_indices == idx)[0][0]

print(f"两种方法结果是否一致: {torch.allclose(unique_indices_optimized.float(), unique_indices_loop.float())}")