按照group取tensor中最大值的下标

1.创建一个tensor，取每个元素对应的组存在index里面，如下例，5个group。输出每个group最大值在对应group中的下标。

输出结果为 

即 每个group中 的最大值 在每个group的下标。 index的值0-4表示5个group，index中元素值相同的表示在 tensor中是一个group。

import torch

# 假设tensor为N*1维的tensor
tensor = torch.tensor([1, 2, 4, 3, 5, 6, 7, 8, 9, 1,41])

# 假设index为N*1维的tensor
index = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 4])

# 将tensor和index合并为一个N*2的tensor
combined_tensor = torch.stack((index, tensor), dim=1)

# 使用torch.max函数取出每一组的最大值的下标
max_indices = torch.zeros(index.max().item() + 1, dtype=torch.long)
for i in range(index.max().item() + 1):
    max_indices[i] = combined_tensor[combined_tensor[:, 0] == i, 1].argmax()

print(max_indices)

2.

1.创建一个tensor，取每个元素对应的组存在index里面，如下例，5个group。输出每个group最大值在对应原始中的下标。

输出结果为 

即 每个group中 的最大值 在每个group的下标。 index的值0-4表示5个group，index中元素值相同的表示在 tensor中是一个group。

import torch

# 假设tensor为N*1维的tensor
tensor = torch.tensor([1, 3, 2, 5, 4, 7, 6, 1,9, 8, 10,41])

# 假设index为N*1维的tensor
index = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 3, 4, 4,4])

max_indices = []
for i in range(index.max().item() + 1):
    group = (index == i)  # 找出属于当前组的元素
    max_index = torch.argmax(tensor[group])  # 找出当前组中的最大值的索引
    true_index = torch.nonzero(group)[max_index].item()  # 找到在原始tensor中的索引
    max_indices.append(true_index)

print(max_indices)  # 输出每一组最大值在原始tensor中的下标