How to retrieve nodes from a sklearn.neighbors.KDTree?

Question:

Is there a way to get a node by id, or all nodes, from a sklearn.neighbors.KDTree instance?

  from sklearn.neighbors import KDTree
  import numpy as np
  
  tree = KDTree(np.array([[0., 0., 0.], [1., 1., 1.]]))
  # How to get the array [[0., 0., 0.], [1., 1., 1.]] back out?
  # Or something like this:
  tree.get_node(0)
  # Which would return: [0., 0., 0.]
Asked By: Robin De Schepper

||

Answers:

The .get_arrays() method gives you access to all arrays:

from sklearn.neighbors import KDTree
import numpy as np

X = np.array([[-2, -2], [-1, -1], [-2, -1], [-3, -2], [0, 0],
              [1, 1], [2, 1], [2, 2], [3, 2]])
kdt = KDTree(X, leaf_size=3, metric='euclidean')
tree_data, index, tree_nodes, node_bounds = kdt.get_arrays()
tree_nodes

Output

array([(0, 9, 0, 3.60555128), (0, 4, 1, 1.11803399),
       (4, 9, 1, 1.80277564)],
      dtype=[('idx_start', '<i8'), ('idx_end', '<i8'), ('is_leaf', '<i8'), ('radius', '<f8')])
Answered By: KRKirov