NumPy：获取插入排序数组后元素的实际索引

请考虑以下代码：

import numpy as np

v = np.linspace(0, 9, 10)
w = np.array([3.5, 4.5])
idx = np.searchsorted(v, w)
v = np.insert(v, idx, w)

print(idx, v[idx])

其输出结果为：

[4 5] array([3.5, 4. ])

变量idx存储的是将w的元素逐个插入v时对应的索引，但当按上述方式将数组整体插入后，只有w最小值对应的idx能正确指向新数组中该元素的位置。

是否存在可用的NumPy函数，能够获取w的元素插入后的实际索引？

没有直接的NumPy函数可以直接返回插入后的实际索引，但可以通过简单的偏移计算得到结果：

情况1：w是已排序数组 #

因为searchsorted要求输入的v是排序状态，若w也保持有序，插入后的实际索引等于原idx加上该元素之前插入的元素个数（即前缀偏移）：

import numpy as np

v = np.linspace(0, 9, 10)
w = np.array([3.5, 4.5])
idx = np.searchsorted(v, w)
# 计算实际索引：原索引 + 前面已插入的元素数量
actual_idx = idx + np.arange(len(idx))
v = np.insert(v, idx, w)

print(actual_idx, v[actual_idx])

输出结果：

[4 6] array([3.5, 4.5])

情况2：w是未排序数组 #

需要先对w和对应的idx排序，计算偏移后再恢复原顺序：

import numpy as np

v = np.linspace(0, 9, 10)
w = np.array([4.5, 3.5])  # 未排序的w
idx = np.searchsorted(v, w)

# 按插入位置排序，获取排序后的索引顺序
sorted_indices = np.argsort(idx)
sorted_idx = idx[sorted_indices]
# 计算排序后的实际索引
sorted_actual_idx = sorted_idx + np.arange(len(sorted_idx))
# 恢复原顺序
actual_idx = np.empty_like(sorted_actual_idx)
actual_idx[sorted_indices] = sorted_actual_idx

v = np.insert(v, idx, w)

print(actual_idx, v[actual_idx])

输出结果：

[6 4] array([4.5, 3.5])

核心逻辑：每插入一个元素，会让后续所有插入位置的索引增加1。对于有序的idx序列，第i个元素的实际索引为原idx[i]加上i；若w无序，则先排序处理偏移，再还原回原元素顺序即可。

内容的提问来源于stack exchange，提问作者G. Gare