Get the distance to each nearest element in a 1D/2D without for loop

Viewed 674

I'm working in python using numpy (could be a pandas series too) and am trying to make the following calculation:

Lets say I have an array corresponding to points on the x axis:

2, 9, 5, 6, 55, 8

For each element in this array I would like to get the distance to the closest element so the output would look like the following:

3, 1, 1, 1, 46, 1

I am trying to find a solution that can scale to 2D (distance to nearest XY point) and ideally would avoid a for loop. Is that possible?

4 Answers

There seems to be a theme with O(N^2) solutions here. For 1D, it's quite simple to get O(N log N):

x = np.array([2, 9, 5, 6, 55, 8])
i = np.argsort(x)
dist = np.diff(x[i])
min_dist = np.r_[dist[0], np.minimum(dist[1:], dist[:-1]), dist[-1]])
min_dist = min_dist[np.argsort(i)]

This clearly won't scale well to multiple dimensions, so use scipy.special.KDTree instead. Assuming your data is N-dimensional and has shape (M, N), you can do

k = KDTree(data)
dist = k.query(data, k=2)[0][:, -1]

Scipy has a Cython implementation of KDTree, cKDTree. Sklearn has a sklearn.neighbors.KDTree with a similar interface as well.

Approach 1

You can use broadcasting in order to get matrix of distances:

>>> data = np.array([2,9,5,6,55,8])
>>> dst_matrix = data - data[:, None]
>>> dst_matrix
array([[  0,   7,   3,   4,  53,   6],
       [ -7,   0,  -4,  -3,  46,  -1],
       [ -3,   4,   0,   1,  50,   3],
       [ -4,   3,  -1,   0,  49,   2],
       [-53, -46, -50, -49,   0, -47],
       [ -6,   1,  -3,  -2,  47,   0]])

Then we can eliminate diagonal as proposed in this post:

dst_matrix = dst_matrix[~np.eye(dst_matrix.shape[0],dtype=bool)].reshape(dst_matrix.shape[0],-1)
>>> dst_matrix
array([[  7,   3,   4,  53,   6],
       [ -7,  -4,  -3,  46,  -1],
       [ -3,   4,   1,  50,   3],
       [ -4,   3,  -1,  49,   2],
       [-53, -46, -50, -49, -47],
       [ -6,   1,  -3,  -2,  47]])

Finally, mininum items can be found:

>>> np.min(np.abs(dst_matrix), axis=1)
array([ 3,  1,  1,  1, 46,  1])

Approach 2 If you're looking for time and memory efficient solution, the best option is scipy.spatial.cKDTrees which packs points (of any dimension) into specific data structure that is optimized for querying closest points. It can also be extended to 2D or 3D.

import scipy.spatial
data = np.array([2,9,5,6,55,8])

ckdtree = scipy.spatial.cKDTree(data[:,None])
distances, idx = ckdtree.query(data[:,None], k=2) 
output = distances[:,1] #distances to not coincident points

For each point querying first two closest points is required here because first of them is expected to be coincident. This is the only solution I found between all the proposed answers that doesn't take ages (the average performance is 4secs for 1M points). Warning: you need to filter duplicated points before applying this method.

There are many ways of achieving it. Some readable and generalizable ways are:

Approach 1:

dist = np.abs(a[:,None]-a)
np.min(dist, where=~np.eye(len(a),dtype=bool), initial=dist.max(), axis=1)
#[ 3  1  1  1 46  1]

Approach 2:

dist = np.abs(np.subtract.outer(a,a))
np.min(dist, where=~np.eye(len(a),dtype=bool), initial=dist.max(), axis=1)

For a 2-D case approach 1 (assumes Euclidean distance. Any other is also possible):

from scipy.spatial.distance import cdist
dist = cdist(a,a)
np.min(dist, where=~np.eye(len(a),dtype=bool), initial=dist.max(), axis=1)

For a 2-D case approach 2 with numpy only:

dist=np.sqrt(((a[:,None]-a)**2).sum(-1))
np.min(dist, where=~np.eye(len(a),dtype=bool), initial=dist.max(), axis=1)

You can achieve a faster distance calculation by using np.dot.

You can do some list comprehension on a pandas series:

s = pd.Series([2,9,5,6,55,8])
s.apply(lambda x: min([abs(x - s[y]) for y in s.index if s[y] != x]))
Out[1]: 
0     3
1     1
2     1
3     1
4    46
5     1

Then you can just add .to_list() or .to_numpy() to the end to get rid of the series index:

s.apply(lambda x: min([abs(x - s[y]) for y in s.index if s[y] != x])).to_numpy()

array([ 3,  1,  1,  1, 46,  1], dtype=int64)
Related