CODEWITHSUNDEEP
pythonkdtree
Dan Schwartz
Dan Schwartz

Reputation: 79

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

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?

Upvotes: 3

Views: 1750

Answers (2)

Jon
Jon

Reputation: 612

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.]]

Upvotes: 3

Dan Schwartz
Dan Schwartz

Reputation: 79

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.

Upvotes: 1

Related Questions