在 NumPy 数组中查找索引的值和最近的 n 个邻居

Question

I wondered if anyone could tell me how to find the index of a number and the indexes of the n-closest neighbors in a NumPy array.

For example, In this array, I would like to find the index of the value 87 , and its four closest neighbors 86 , 88 to the left and 78 , 43 to the right.

a = np.random.randint(1,101,15)
array([79, 86, 88, 87, 78, 43, 57])

Answer 1

If you want to change the values from time to time, although this would be expensive for large arrays, should do the trick:

a = np.array([79, 86, 88, 87, 78, 43, 57])

number = 87
n_nearest = 4

index = np.where(a == number)[0][0] # np.where returns (array([3]),)

a = a[max(0, index - n_nearest // 2):index + n_nearest // 2 + 1]
nearests = np.delete(a, n_nearest // 2)

print(nearests)

Output: [86 88 78 43]

First, find the index of the value you get the neighbors (may not work with duplicate values though).

You should do max(0, index - 2) in case of the value you want may be at the beginning of array (position 0 or 1).

Then, delete the number from the result. The rest of it will be the neighbors you want.

Answer 2

I had a go at it, with the caveat that I'm not hugely experienced with python or numpy - only a couple of months

(...so I'd also look for someone else to chip in a much cleaner/simpler/better method!)

from functools import reduce
import operator

a = np.array([5, 10, 15, 12, 88, 86, 5, 87, 1,2,3, 87,1,2,3])

look_for = 87

# find indicies where a == 87:
np.nonzero(a == look_for)

# get as interable
np.nonzero(a == look_for)[0]

# put into list comprehension with the delta in indicies you want and the values
# from above inside 'range' to generate a range of values b/w index-delta:index+delta,
# then wrap it into a list to generate the list from the range iterator:
delta = 2
[list(range(i-delta,i+delta+1)) for i in np.nonzero(a==87)[0]]

# above gives you a list of lists, need to flatten into a single list
reduce(operator.concat, [list(range(i-delta,i+delta+1)) for i in np.nonzero(a==87)[0]])

# want only unique values, so one way is to turn above into a set
set(reduce(operator.concat, [list(range(i-delta,i+delta+1)) for i in np.nonzero(a==87)[0]]))

# the above gives you a set with all the indicies, with only unique values.
# one remaning problem is it still could have values < 0 or > a.size, so
# you'd now want to wrap it all into another list comprehension to get rid of 
# any values < 0 or > a.size