This is similar to using bisect_left, but it'll allow you to pass in an array of targets
def find_closest(A, target):
#A must be sorted
idx = A.searchsorted(target)
idx = np.clip(idx, 1, len(A)-1)
left = A[idx-1]
right = A[idx]
idx -= target - left < right - target
return idx
Some explanation:
First the general case: idx = A.searchsorted(target)
returns an index for each target
such that target
is between A[index - 1]
and A[index]
. I call these left
and right
so we know that left < target <= right
. target - left < right - target
is True
(or 1) when target is closer to left
and False
(or 0) when target is closer to right
.
Now the special case: when target
is less than all the elements of A
, idx = 0
. idx = np.clip(idx, 1, len(A)-1)
replaces all values of idx
< 1 with 1, so idx=1
. In this case left = A[0]
, right = A[1]
and we know that target <= left <= right
. Therefor we know that target - left <= 0
and right - target >= 0
so target - left < right - target
is True
unless target == left == right
and idx - True = 0
.
There is another special case if target
is greater than all the elements of A
, In that case idx = A.searchsorted(target)
and np.clip(idx, 1, len(A)-1)
replaces len(A)
with len(A) - 1
so idx=len(A) -1
and target - left < right - target
ends up False
so idx returns len(A) -1
. I'll let you work though the logic on your own.
For example:
In [163]: A = np.arange(0, 20.)
In [164]: target = np.array([-2, 100., 2., 2.4, 2.5, 2.6])
In [165]: find_closest(A, target)
Out[165]: array([ 0, 19, 2, 2, 3, 3])
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…