def query(self, query_points, k):
        """Query the tree

        Parameters
        ----------
        query_points: numpy.ndarray (of type `np.float32`)
            Data points for which the next neighbours should be calculated
        k: int
            Number of neighbors to find

        Returns
        -------
        distances: numpy.ndarray
            Distances to the neighbors
        indices: numpy.ndarray
            Indices of the k nearest neighbors in original data array

        Example
        -------
        >>> from pymatting import *
        >>> import numpy as np
        >>> data_set = np.random.randn(100, 2)
        >>> tree = KDTree(data_set.astype(np.float32))
        >>> tree.query(np.array([[0.5,0.5]], dtype=np.float32), k=3)
        (array([[0.14234178, 0.15879704, 0.26760164]], dtype=float32), array([[29, 21, 20]]))
        """
        assert query_points.dtype == np.float32

        if self.shuffled_data_points.shape[0] < k:
            raise ValueError(
                f"Number of data points ({self.shuffled_data_points.shape[0]}) is less than the number of neighbors requested (k={k})."
                " Please provide a larger dataset or reduce the value of k."
            )

        n_query = query_points.shape[0]

        squared_distances = np.empty((n_query, k), np.float32)
        indices = np.empty((n_query, k), np.int64)

        _find_knn(
            self.i0_inds,
            self.i1_inds,
            self.less_inds,
            self.more_inds,
            self.split_dims,
            self.bounds,
            self.split_values,
            self.shuffled_data_points,
            query_points,
            indices,
            squared_distances,
            k,
        )

        indices = self.shuffled_indices[indices]
        distances = np.sqrt(squared_distances)

        return distances, indices
Query the tree

        Parameters
        ----------
        query_points: numpy.ndarray (of type `np.float32`)
            Data points for which the next neighbours should be calculated
        k: int
            Number of neighbors to find

        Returns
        -------
        distances: numpy.ndarray
            Distances to the neighbors
        indices: numpy.ndarray
            Indices of the k nearest neighbors in original data array

        Example
        -------
        >>> from pymatting import *
        >>> import numpy as np
        >>> data_set = np.random.randn(100, 2)
        >>> tree = KDTree(data_set.astype(np.float32))
        >>> tree.query(np.array([[0.5,0.5]], dtype=np.float32), k=3)
        (array([[0.14234178, 0.15879704, 0.26760164]], dtype=float32), array([[29, 21, 20]]))
        
pymatting/util/kdtree.py