KDTree building optimisation: split algorithm changed (#210)

Author: gyrdym

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 16.6.2
+- KDTree:
+    - KDTree build optimization: split algorithm changed
+
 ## 16.6.1
 - `KDTree` class added to library export file
 

benchmark/kd_tree_building.dart

@@ -1,4 +1,4 @@
-// 0.5 sec (MacBook Air mid 2017)
+// 0.8 sec (MacBook Air mid 2017)
 import 'dart:convert';
 import 'dart:io';
 

benchmark/kd_tree_querying.dart

@@ -1,4 +1,4 @@
-// 0.04 sec (MacBook Air mid 2017)
+// 0.03 sec (MacBook Air mid 2017)
 import 'dart:convert';
 import 'dart:io';
 
@@ -35,7 +35,7 @@ Future main() async {
   final decodedPoints = jsonDecode(dataAsString) as Map<String, dynamic>;
 
   trainData = DataFrame.fromJson(decodedPoints);
-  tree = KDTree(trainData);
+  tree = KDTree(trainData, leafSie: 1);
   point = Vector.randomFilled(trainData.rows.first.length,
       seed: 10, min: -5000, max: 5000);
 

lib/src/retrieval/kd_tree/kd_tree.dart

@@ -40,13 +40,16 @@ abstract class KDTree implements Serializable {
   /// the source [points] matrix. Example:
   ///
   /// ```dart
+  /// import 'package:ml_dataframe/ml_dataframe.dart';
+  /// import 'package:ml_linalg/vector.dart';
+  ///
   /// final data = DataFrame([
   ///   [21, 34, 22, 11],
   ///   [11, 33, 44, 55],
   ///   ...,
   /// ], headerExists: false);
   /// final kdTree = KDTree(data);
-  /// final neighbours = kdTree.query([1, 2, 3, 4], 2);
+  /// final neighbours = kdTree.query(Vector.fromList([1, 2, 3, 4]), 2);
   ///
   /// print(neighbours.index); // let's say, it outputs `3` which means that the nearest neighbour is kdTree.points[3]
   /// ```

lib/src/retrieval/kd_tree/kd_tree_builder.dart

@@ -19,15 +19,14 @@ class KDTreeBuilder {
 
   KDTreeNode _train(List<int> pointIndices) {
     final isLeaf = pointIndices.length <= _leafSize;
+    final points = _points.sample(rowIndices: pointIndices);
+    final splitIdx = _getSplitIdx(points);
 
     if (isLeaf) {
-      return KDTreeNode(pointIndices: pointIndices);
+      return KDTreeNode(splitIndex: splitIdx, pointIndices: pointIndices);
     }
 
-    final points = _points.sample(rowIndices: pointIndices);
-    final splitIdx = _getSplitIdx(points);
-    final splitValue = points.getColumn(splitIdx).median();
-    final split = _splitPoints(pointIndices, splitIdx, splitValue);
+    final split = _splitPoints(pointIndices, splitIdx);
 
     return KDTreeNode(
       pointIndices: [split.midPoint],
@@ -55,27 +54,15 @@ class KDTreeBuilder {
     return maxIdx;
   }
 
-  _Split _splitPoints(List<int> pointIndices, int splitIdx, num splitValue) {
-    final left = <int>[];
-    final right = <int>[];
-    int? midPoint;
-
-    for (var i = 0; i < pointIndices.length; i++) {
-      final pointIndex = pointIndices[i];
-      final point = _points[pointIndex];
+  _Split _splitPoints(List<int> pointIndices, int splitIdx) {
+    pointIndices.sort((firstIdx, secondIdx) =>
+        _points[firstIdx][splitIdx].compareTo(_points[secondIdx][splitIdx]));
 
-      if (point[splitIdx] < splitValue) {
-        left.add(pointIndex);
-        continue;
-      }
-
-      if (midPoint == null || point[splitIdx] < _points[midPoint][splitIdx]) {
-        midPoint = pointIndex;
-      }
-
-      right.add(pointIndex);
-    }
+    final midPointIdx = (pointIndices.length / 2).floor();
+    final midPoint = pointIndices[midPointIdx];
+    final left = pointIndices.sublist(0, midPointIdx);
+    final right = pointIndices.sublist(midPointIdx + 1);
 
-    return _Split(left, right, midPoint!);
+    return _Split(left, right, midPoint);
   }
 }

lib/src/retrieval/kd_tree/kd_tree_impl.dart

@@ -61,9 +61,11 @@ class KDTreeImpl with SerializableMixin implements KDTree {
     return neighbours.toList().reversed;
   }
 
-  void _findKNNRecursively(KDTreeNode node, Vector point, int k,
+  void _findKNNRecursively(KDTreeNode? node, Vector point, int k,
       HeapPriorityQueue<KDTreeNeighbour> neighbours) {
-    searchIterationCount++;
+    if (node == null) {
+      return;
+    }
 
     if (node.isLeaf) {
       _knnSearch(point, node.pointIndices, neighbours, k);
@@ -72,31 +74,30 @@ class KDTreeImpl with SerializableMixin implements KDTree {
     }
 
     final nodePoint = points[node.pointIndices[0]];
-    final boundaryPoint = _getBoundary(point, node);
     final isNodeTooFar = neighbours.length > 0 &&
-        boundaryPoint.distanceTo(nodePoint) > neighbours.first.distance;
+        (point[node.splitIndex] - nodePoint[node.splitIndex]).abs() >
+            neighbours.first.distance;
     final isQueueFilled = neighbours.length == k;
 
     if (isQueueFilled && isNodeTooFar) {
       return;
     }
 
-    if (node.left == null) {
-      _findKNNRecursively(node.right!, point, k, neighbours);
-    } else if (node.right == null) {
-      _findKNNRecursively(node.left!, point, k, neighbours);
-    } else if (point[node.splitIndex!] < nodePoint[node.splitIndex!]) {
-      _findKNNRecursively(node.left!, point, k, neighbours);
-      _findKNNRecursively(node.right!, point, k, neighbours);
+    _knnSearch(point, node.pointIndices, neighbours, k);
+
+    if (point[node.splitIndex] < nodePoint[node.splitIndex]) {
+      _findKNNRecursively(node.left, point, k, neighbours);
+      _findKNNRecursively(node.right, point, k, neighbours);
     } else {
-      _findKNNRecursively(node.right!, point, k, neighbours);
-      _findKNNRecursively(node.left!, point, k, neighbours);
+      _findKNNRecursively(node.right, point, k, neighbours);
+      _findKNNRecursively(node.left, point, k, neighbours);
     }
   }
 
   void _knnSearch(Vector point, List<int> pointIndices,
       HeapPriorityQueue<KDTreeNeighbour> neighbours, int k) {
     pointIndices.forEach((candidateIdx) {
+      searchIterationCount++;
       final candidate = points[candidateIdx];
       final candidateDistance = candidate.distanceTo(point);
       final lastNeighbourDistance =
@@ -113,13 +114,4 @@ class KDTreeImpl with SerializableMixin implements KDTree {
       }
     });
   }
-
-  Vector _getBoundary(Vector point, KDTreeNode node) {
-    final nodePoint = points[node.pointIndices[0]];
-    final boundarySrc = [...nodePoint];
-
-    boundarySrc[node.splitIndex!] = point[node.splitIndex!];
-
-    return Vector.fromList(boundarySrc, dtype: point.dtype);
-  }
 }

lib/src/retrieval/kd_tree/kd_tree_neighbour.dart

@@ -14,5 +14,5 @@ class KDTreeNeighbour {
   }
 
   @override
-  String toString() => 'Index: $index, Distance: $distance';
+  String toString() => '(Index: $index, Distance: $distance)';
 }

lib/src/retrieval/kd_tree/kd_tree_node.dart

@@ -7,15 +7,18 @@ part 'kd_tree_node.g.dart';
 @JsonSerializable()
 class KDTreeNode {
   KDTreeNode(
-      {this.splitIndex, this.left, this.right, required this.pointIndices});
+      {required this.splitIndex,
+      this.left,
+      this.right,
+      required this.pointIndices});
 
   factory KDTreeNode.fromJson(Map<String, dynamic> json) =>
       _$KDTreeNodeFromJson(json);
 
   Map<String, dynamic> toJson() => _$KDTreeNodeToJson(this);
 
   @JsonKey(name: kdTreeNodeIndexJsonKey)
-  final int? splitIndex;
+  final int splitIndex;
 
   @JsonKey(name: kdTreeNodeLeftJsonKey)
   final KDTreeNode? left;

lib/src/retrieval/kd_tree/kd_tree_node.g.dart

@@ -1,6 +1,6 @@
 name: ml_algo
 description: Machine learning algorithms, Machine learning models performance evaluation functionality
-version: 16.6.1
+version: 16.6.2
 homepage: https://github.com/gyrdym/ml_algo
 
 environment: