Add KdTree size and depth methods (#603)

Signed-off-by: Martin Davis <[email protected]>
locationtech · Sep 25, 2020 · 4867e71 · 4867e71
1 parent 5c78e72
commit 4867e71
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 1 deletion.
diff --git a/modules/core/src/main/java/org/locationtech/jts/index/kdtree/KdTree.java b/modules/core/src/main/java/org/locationtech/jts/index/kdtree/KdTree.java
@@ -253,14 +253,16 @@ private KdNode insertExact(Coordinate p, Object data) {
       }
       leafNode = currentNode;
       if (isLessThan) {
+        //System.out.print("L");
         currentNode = currentNode.getLeft();
       } else {
+        //System.out.print("R");
         currentNode = currentNode.getRight();
       }
 
       isOddLevel = ! isOddLevel;
     }
-
+    //System.out.println("<<");
     // no node found, add new leaf node to tree
     numberOfNodes = numberOfNodes + 1;
     KdNode node = new KdNode(p, data);
@@ -383,4 +385,39 @@ public KdNode query(Coordinate queryPt) {
     return queryNodePoint(root, queryPt, true);
   }
 
+  /**
+   * Computes the depth of the tree.
+   * 
+   * @return the depth of the tree
+   */
+  public int depth() {
+    return depthNode(root);
+  }
+
+  private int depthNode(KdNode currentNode) {
+    if (currentNode == null)
+      return 0;
+
+    int dL = depthNode(currentNode.getLeft());
+    int dR = depthNode(currentNode.getRight());
+    return 1 + (dL > dR ? dL : dR);
+  }
+
+  /**
+   * Computes the size (number of items) in the tree.
+   * 
+   * @return the size of the tree
+   */
+  public int size() {
+    return sizeNode(root);
+  }
+
+  private int sizeNode(KdNode currentNode) {
+    if (currentNode == null)
+      return 0;
+
+    int sizeL = sizeNode(currentNode.getLeft());
+    int sizeR = sizeNode(currentNode.getRight());
+    return 1 + sizeL + sizeR;
+  }
 }
diff --git a/modules/core/src/test/java/org/locationtech/jts/index/kdtree/KdTreeTest.java b/modules/core/src/test/java/org/locationtech/jts/index/kdtree/KdTreeTest.java
@@ -92,6 +92,17 @@ public void testSnapToNearest() {
         "MULTIPOINT ( (10 60), (20 60), (20 60))");
   }
 
+  public void testSizeDepth() {
+    KdTree index = build("MULTIPOINT ( (10 60), (20 60), (16 60), (1 1), (23 400))", 
+        0);
+    int size = index.size();
+    assertEquals(5, size);
+    int depth = index.depth();
+    // these are weak conditions, but depth varies depending on data and algorithm
+    assertTrue( depth > 1 );
+    assertTrue( depth <= size );
+  }
+
   private void testQuery(String wktInput, double tolerance,
       Envelope queryEnv, String wktExpected) {
     KdTree index = build(wktInput, tolerance);