查找接近目标的所有值，像numpy.searchsorted()但返回所有相同值？

Question

有没有什么好的方法可以在一个已排序的数组A中找到所有接近几个目标值的索引呢？使用numpy的searchsorted()可以让我们高效地找到接近多个目标的索引：在Python中找到最近的值并返回数组的索引。但是，如果数组A中有重复的值，这种方法只会返回其中一个索引，而不是所有可能的索引。

比如说，像这样的数组：

A = array([    1. ,     2. ,     3. ,     3. ,     3.1,     4. ,    50. ,
          60. ,    70. ,    80. ,    90. ,   100.1,   110. ,   120. ,
         999. ,  1000. ])
targets=[3, 100]

它会返回idx = [2, 11]，但我希望它返回[[2,3],11]。我可以通过循环idx来获取布尔索引，比如[A==A[idx[0]],A==A[idx[1]],...]，但是如果目标数组非常大，这样做会非常低效。

我可以先用numpy.unique()找到数组中的唯一值集合，以找到所有相同的值。然后在这个唯一数组上使用searchsorted()，这样可能会节省一些时间。接着，我可以用这个索引找到所有相同的值。

这里有一个例子：

def find_closest_multiTargets_inSortred(A,targets):
        #A must be sorted
    idx = A.searchsorted(targets)
    idx = npy.clip(idx, 1, len(A)-1)
    left = A[idx-1]
    right = A[idx]
    idx -= targets - left < right - targets
    return idx

def find_closest_multiTargets_Allrepeats(A,targets):
    ua=npy.unique(A)
    _uaIdxs=find_closest_multiTargets_inSortred(ua, targets)
    return [npy.where(A==ua[_i]) for _i in _uaIdxs]

>>> find_closest_multiTargets_Allrepeats([5.1,5.5,4,1,2.3,5.1,6],[2,5])
[(array([4]),), (array([0, 5]),)]

我认为，如果len(ua)<<len(A)，那么在A上直接寻找最近的值会效率低很多。然而，npy.where这一步还是要循环遍历_uaIdxs，如果它很大，那就会非常低效。如果能构建一个替代的unique()，为A中的每个唯一值获取一个索引列表（[[索引有值ua[0]],[索引有值ua[2]]...]），那将会高效得多：

def find_closest_multiTargets_Allrepeats2(A,targets):
    ua,idxList=npy.unique2(A)
    _uaIdxs=find_closest_multiTargets_inSortred(ua, targets)
    return idxList[_uaIdxs]

但我不知道是否有任何方法可以实现unique2()所期望的功能。可能还有其他完全不同的算法可以以更高效的方式获得相同的结果，而不仅仅是searchsorted。

为了简单起见，我们假设A是已排序的。对于未排序的数组A，我们可以先对它进行排序。

有没有人能提供一个更高效的方法来做到这一点？

谢谢！

Answer 1

numpy.in1d(A, idx) 这个函数可以实现你想要的功能。

Answer 2

你可以这样做：

a = np.array([1., 2., 3., 3., 3.1, 4., 50., 60., 70., 80., 90., 100.1, 110., 120., 999., 1000.])
t = np.array([3, 100])

• 计算每一对之间的距离：

  d = np.abs(np.subtract.outer(a, t))
  

• 找到最近的值：

  asort = np.argsort(d, axis=0)
  

• 获取最近的索引和最近的值：

  ind = np.arange(a.shape[0])
  print(ind[asort][0])
  #array([ 2, 11], dtype=int64)
  
  print(a[asort][0])
  #array([   3. ,  100.1])
  

注意，如果在最后一步使用其他索引 [i]，而不是 [0]，你会得到第 i 个最近的值……使用 [0] 会得到最近的值。