TensorFlow中如何获取向量的Sorted Index Function（排序索引函数）？

嘿，这个问题我刚好处理过！你要的排序索引函数ψ，其实就是能让排序后的向量等于原向量通过该索引张量取值的那个ψ对吧？简单来说，假设原向量是v，排序后的向量是sorted_v，那么ψ满足sorted_v[i] = v[ψ[i]]。在TensorFlow里，用tf.argsort就能直接获取这个索引张量，我给你一步步演示：

基础用法（一维向量） #

首先，我们先创建一个示例向量，然后用tf.argsort生成ψ：

import tensorflow as tf

# 定义原向量v
v = tf.constant([3, 1, 4, 1, 5])

# 获取升序排序的索引ψ
psi = tf.argsort(v)

验证一下是否符合预期：

# 通过ψ索引原向量，得到排序后的结果
sorted_v = tf.gather(v, psi)

print("原向量:", v.numpy())
print("排序索引ψ:", psi.numpy())
print("排序后的向量:", sorted_v.numpy())

输出结果：

原向量: [3 1 4 1 5]
排序索引ψ: [1 3 0 2 4]
排序后的向量: [1 1 3 4 5]

降序排序的ψ #

如果需要降序排序对应的索引，只需要给tf.argsort加上direction='DESCENDING'参数：

psi_desc = tf.argsort(v, direction='DESCENDING')
sorted_v_desc = tf.gather(v, psi_desc)

print("降序排序索引ψ:", psi_desc.numpy())
print("降序排序后的向量:", sorted_v_desc.numpy())

输出结果：

降序排序索引ψ: [4 2 0 1 3]
降序排序后的向量: [5 4 3 1 1]

高维张量的情况 #

如果你的张量是高维的（比如二维矩阵），可以通过axis参数指定按哪一维度排序获取索引。比如按行排序：

v_2d = tf.constant([[3, 1, 2], [5, 0, 4]])
# 按行（axis=1）获取排序索引ψ
psi_2d_row = tf.argsort(v_2d, axis=1)

print("二维张量按行排序的ψ:\n", psi_2d_row.numpy())

输出结果：

二维张量按行排序的ψ:
 [[1 2 0]
 [1 2 0]]

简单来说，tf.argsort返回的张量就是你要找的排序索引函数ψ，它直接给出了排序后每个位置对应的原向量元素索引，完全符合你的需求～

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