np.searchsorted与np.argmin查找最近索引结果差异原因咨询

我有一组时间戳数组arr，还有包含多组起止时间的列表cuts，想要截取每个起止时间之间的时间戳生成新数组。但用np.searchsorted()和np.argmin()两种方法得到的索引不一样，想知道原因。

我的代码 #

import numpy as np

# Initialization data  
arr = np.arange(761.55643, 1525.5704932002686, 1/ 1000)

cuts = [[810.211186646, 899.102014549], [903.520741867, 982.000921478], [985.201032795, 993.400610844], 
       [998.303881868, 1085.500698357], [1090.200656211, 1168.101925871], [1171.299249968, 1179.611318749], 
       [1184.610645285, 1271.597569677], [1275.600586067, 1363.696138556], [1368.301122947, 1455.500707533]]

vector_validity = np.zeros(len(arr))
new_arr_with_argmin = np.zeros(0)
for cut in cuts:
    vector_validity[int(np.searchsorted(arr, cut[0])) : int(np.searchsorted(arr, cut[1]))] = 1
    print(f"np.searchsorted start: {np.searchsorted(arr, cut[0])}")
    print(f"np.argmin start: {np.argmin(abs(arr - cut[0]))}")
    print(f"np.searchsorted end: {np.searchsorted(arr, cut[1])}")
    print(f"np.argmin end: {np.argmin(abs(arr - cut[1]))}")
    
    new_arr_with_argmin = np.concatenate((new_arr_with_argmin, arr[np.argmin(abs(arr - cut[0])) : np.argmin(abs(arr - cut[1]))]))
new_arr_with_searchsorted = arr[vector_validity == 1] 

打印结果 #

np.searchsorted start: 48655
np.argmin start: 48655
np.searchsorted end: 137546
np.argmin end: 137546
np.searchsorted start: 141965
np.argmin start: 141964
np.searchsorted end: 220445
np.argmin end: 220444
np.searchsorted start: 223645
np.argmin start: 223645
np.searchsorted end: 231845
np.argmin end: 231844
np.searchsorted start: 236748
np.argmin start: 236747
np.searchsorted end: 323945
np.argmin end: 323944
np.searchsorted start: 328645
np.argmin start: 328644
np.searchsorted end: 406546
np.argmin end: 406545
np.searchsorted start: 409743
np.argmin start: 409743
np.searchsorted end: 418055
np.argmin end: 418055
np.searchsorted start: 423055
np.argmin start: 423054
np.searchsorted end: 510042
np.argmin end: 510041
np.searchsorted start: 514045
np.argmin start: 514044
np.searchsorted end: 602140
np.argmin end: 602140
np.searchsorted start: 606745
np.argmin start: 606745
np.searchsorted end: 693945
np.argmin end: 693944

原因解释 #

两个函数的设计逻辑完全不同，这是导致结果差异的核心：

• np.searchsorted()：专门针对有序数组，返回的是目标值应该插入的位置，保证插入后数组仍保持有序。默认使用side='left'参数，意味着当目标值不在数组中时，返回第一个大于目标值的元素索引；如果目标值存在，则返回第一个等于目标值的元素索引。
• np.argmin(abs(arr - x))：计算数组中每个元素与目标值的绝对差，返回绝对差最小的元素索引，也就是找数组里最接近目标值的元素位置。

结合你的数据来看：
你的arr是步长为0.001的有序序列，但cuts中的起止时间大多不是arr中的精确值。当目标值落在arr[i]和arr[i+1]之间时：

• np.searchsorted()会返回i+1，因为要插入到这个位置才能维持数组的有序性。
• np.argmin()则会比较目标值到arr[i]和arr[i+1]的距离，返回更近的那个元素的索引。比如第二个区间的起始值903.520741867，arr[141964] = 903.52043，arr[141965] = 903.52143，目标值离前者更近，所以返回141964，而searchsorted返回141965。

当目标值刚好等于arr中的某个元素时，两个函数会返回相同的索引，比如第一个区间的起止时间。

如何选择方法？ #

• 如果需要严格截取[start, end)区间内的所有元素（即大于等于start、小于end的元素），np.searchsorted()更合适，它的逻辑完全匹配有序数组的区间划分。
• 如果需要截取最接近start和end的元素之间的区间，可以用np.argmin()，但要注意这种情况下可能会因为目标值的位置，导致截取的区间比预期多/少一个元素。

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