update

qzhao19 · Sep 8, 2023 · e6e2d5a · e6e2d5a
1 parent 7995684
commit e6e2d5a
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/kdtree/kd_tree.py b/kdtree/kd_tree.py
@@ -204,4 +204,17 @@ def _query_radius_single_data(self, data, radius):
 
                 if self._check_intersection(right_hyper_rect, data, radius):
                     stack.append(self.tree[right])              
-        return inside
+        return inside
+
+    def query_radius(self, data, radius):
+        num_samples, _ = data.shape
+        # search kdtree
+        knn = []
+        for i in np.arange(num_samples):
+            # _data = data[:,i].reshape((param,1)).repeat(leafsize, axis=1)
+            _data = data[i, :][np.newaxis].repeat(self.leaf_size, axis=0)
+            # print(_data)
+            _knn = self._query_radius_single_data(_data, radius)
+            knn.append(_knn)
+
+        return knn