我有一个向量 A = [0,1,2,3,0,0,1,1,2,2,3,3]。我需要以递增的方式对其进行排序，使其按顺序排列，并从中提取argsort。为了更好地解释这一点，我需要对A进行排序，使其返回 B = [0,4,5,1,6,7,2,8,9,3,10,11]。然而，当我使用pytorch的 torch.argsort(A) 时，它返回 B = [4,5,0,1,6,7,2,8,9,3,10,11]。

我假设执行此操作的算法无法由我控制。有什么方法可以在不引入for循环的情况下解决这个问题吗？这样的操作是我神经网络模型的一部分，如果不高效处理会导致性能问题。谢谢！

回答：

这里是一个基于纯PyTorch的解决方案，利用了广播、torch.unique()和torch.nonzero()。这对于基于GPU的实现/运行将是一个巨大的提升，如果我们必须切换回NumPy，使用argsort然后再传回PyTorch（如其他方法所建议的），这是不可能的。

# 我们的输入张量In [50]: A = torch.tensor([0,1,2,3,0,0,1,1,2,2,3,3])# 构造一个中间布尔张量In [51]: boolean = A[:, None] == torch.unique(A)In [52]: booleanOut[52]: tensor([[1, 0, 0, 0],        [0, 1, 0, 0],        [0, 0, 1, 0],        [0, 0, 0, 1],        [1, 0, 0, 0],        [1, 0, 0, 0],        [0, 1, 0, 0],        [0, 1, 0, 0],        [0, 0, 1, 0],        [0, 0, 1, 0],        [0, 0, 0, 1],        [0, 0, 0, 1]], dtype=torch.uint8)

一旦我们有了这个布尔张量，我们可以通过检查转置后的布尔张量中位置为1的地方来找到所需的索引。

这将为我们提供排序后的input和indices。由于我们只想要索引，我们可以通过索引最后一列（1或-1）来获取它们。

In [53]: torch.nonzero(boolean.t())[:, -1]Out[53]: tensor([ 0,  4,  5,  1,  6,  7,  2,  8,  9,  3, 10, 11])

这是原始帖中评论区提供的另一个示例的结果：

In [55]: A_large = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9])In [56]: boolean_large = A_large[:, None] == torch.unique(A_large)In [57]: torch.nonzero(boolean_large.t())[:, -1]Out[57]: tensor([ 0, 10, 11,  1, 12, 13,  2, 14, 15,  3, 16, 17,  4, 18, 19,  5, 20, 21,         6, 22, 23,  7, 24, 25,  8, 26, 27,  9, 28, 29])

注意：与其他答案中提出的基于NumPy的解决方案不同，在这里我们不必担心要使用哪种排序算法，因为我们根本没有使用任何排序。