Get node indices in NearestNeighbors?

Hey all—

I’m building a KDTree using NearestNeighbors.jl like this:

using StaticArrays, NearestNeighbors
pts  = [SVector{2, Float64}(randn(2)) for _ in 1:1024]
tree = KDTree(pts, leafsize=25) 

and I’m trying to find the correct way to get information about what indices of are in leaf node j. In this example there are 41 nodes, and I’d love to have some vector leaf_indices such that leaf_indices[j] is a Vector{Int64} of length 25, or of course some other representation of the same information. I think that indices 1:25 are leaf node 1, 26:50 are in leaf node 2, and so on, but I’m not completely certain.

If that’s not right: I see that there are some internal functions for getting the indices of the (re-ordered) data based on the index of a leaf node, but I’m having trouble figuring out what the indices for the leaf nodes are. NearestNeighbors.isleaf just checks if an index is larger than the number of non-leaf nodes (which in this case is 40), but there’s something I’m not understanding because NearestNeighbors.get_leaf_range(tree.tree_data, 41) is 450:474, and the ranges only increase as I increase the argument. No integer argument here gives me 1:25, for example. So I’m clearly not understanding the design here.

Can anybody help me understand the right way to extract this information?

I’m not sure how bad of manners this is to do…but @kristoffer.carlsson, can you help me out with this, assuming that it is as simple as I think it is?

1 Like

It was quite a long time since I worked properly with the internals of NearestNeighbors so my memory is a bit fuzzy on the details. But perhaps this notebook can help It’s a very old notebook though so it probably doesn’t still run, I should update it.

Anyway, it shows some illustrations for a Ball Tree but I think it should be pretty much the same for a KDTree. For example:

# Skip non leaf nodes
offset = tree.tree_data.n_internal_nodes + 1
nleafs = tree.tree_data.n_leafs

# Range of leaf nodes
index_range = offset: offset + nleafs - 1

leaf_indices = map(idx -> NearestNeighbors.get_leaf_range(tree.tree_data, idx), index_range)
41-element Vector{UnitRange{Int64}}:

This should give you the indices into the reordered data. If you want the indices into your original input data, it should be:

julia> [tree.indices[i] for i in leaf_indices]
41-element Vector{Vector{Int64}}:
 [513, 110, 792, 461, 670, 143, 198, 842, 246, 501  …  744, 945, 341, 807, 862, 825, 739, 26, 677, 96]
 [291, 680, 161, 882, 132, 36, 398, 1010, 727, 345  …  126, 65, 89, 510, 849, 195, 617, 471, 448, 381]
 [169, 783, 772, 477, 355, 323, 614, 684, 202, 768  …  750, 437, 1001, 780, 310, 328, 597, 482, 701, 681]


They will increase until you reach tree.tree_data.cross_node

julia> tree.tree_data.cross_node

julia> NearestNeighbors.get_leaf_range(tree.tree_data, 63)

julia> NearestNeighbors.get_leaf_range(tree.tree_data, 64)

julia> NearestNeighbors.get_leaf_range(tree.tree_data, 65)
1 Like

To give an example, let’s say we create a tree with 6 points and leafsize=1 for simplicity:

Note that in NearestNeighbors.jl, we always fill “from the left” and at most one leaf node will not have a completely filled “bucket”.

The black numbers are the numer of the nodes. The green ones are the number of the leaf nodes. The “cross node” is 8 which is the node where the leaf nodes start (but note that the node number for leaf node 1 is higher than the node number for leaf node 5. Doing this in the package:

julia> tree = KDTree(rand(3,6); leafsize=1);

# nodes 1-5 are internal (non-leaf)
julia> tree.tree_data.n_internal_nodes

# 8 is the "cross" node where the leaf nodes start
julia> tree.tree_data.cross_node

# The range of the leaf nodes for node 8
julia> NearestNeighbors.get_leaf_range(tree.tree_data, 8)

# The range of leaf nodes for node 6
julia> NearestNeighbors.get_leaf_range(tree.tree_data, 6)

# and the size of the last (possible unfilled) node bucket is here of course 1
julia> tree.tree_data.last_node_size

You can try create a tree with 11 nodes and a leafsize of 2 to see how the last bucket will now be half full.


@kristoffer.carlsson: Wow! I really can’t thank you enough. That’s incredibly helpful and definitely enough to work it out exactly for the KD-tree case. I’ll update this thread with a clean solution shortly.

Thanks again for the amazing response, and thanks for writing one of my favorite packages in the Julia ecosystem.

1 Like