How can I retrieve the coordinates of a point in a kdtree given that point's tree index?

user_8202914 0
Viewed 406

For example, suppose I have created a kdtree from the points in a 2-d array. Then I run a nearest neighbor query on some point (represented as a pair of x-y coordinates) asking for the 8 nearest neighbors to that point. This returns two lists: (1) a list of 8 distances to neighbors of the given pint, and (2) a list of tree indices for those neighbors.

Suppose now that I want to take one of those tree indices and get the x-y coordinates of that point? How can I do this?

kdtree python
2 Answers

FYI there is a data attribute that contains the coordinates of the tree: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.html

So you can get your coordinates via

tree.data[index]

Here's a full example:

import numpy as np
from scipy.spatial import KDTree

pts = np.array([(1, 1), (2, 1), (3, 1), (4, 1), (1, 2), (2, 2), (3, 2), (4, 2), (1, 3), (2, 3), (3, 3), (4, 3), (1, 4), (2, 4), (3, 4), (4, 4)])

# Build tree
tree = KDTree(pts)

# Query tree for points within distance=2 of (1,1)
idx = tree.query_ball_point([1,1],r=2)
print(idx)

which yields

[0, 4, 8, 1, 5, 2]

Then to get the coordinates of these indicies,

# Get coordinates of the indices
print(tree.data[idx])
[[1. 1.]
 [1. 2.]
 [1. 3.]
 [2. 1.]
 [2. 2.]
 [3. 1.]]
user_81955280

In case anyone is interested, one of my students provided the answer. There is no kdtree method that returns this information, but you can get it from the original lists used to create the tree. The x-y coordinates of the points used to construct the tree are taken from two lists, one of the x coordinates, and one of the y-coordinates. Then, when you do a kdtree query, the list of indexes returned by the query can be used as indexes into the two original lists to get back the original x-y pairs.

user_82029140
Related