Track cluster sizes of the condensed tree

Co-authored-by: Ioannis Panagiotas <ioannis.panagiotas@neotechnology.com>

Author: vnickolov

algo/src/main/java/org/neo4j/gds/hdbscan/CondenseStep.java

@@ -50,6 +50,8 @@ CondensedTree condense(ClusterHierarchy clusterHierarchy, long minClusterSize) {
         var clusterHierarchyRoot = clusterHierarchy.root();
         var parent = HugeLongArray.newArray(clusterHierarchyRoot + 1);
         var lambda = HugeDoubleArray.newArray(clusterHierarchyRoot + 1);
+        var size = HugeLongArray.newArray(nodeCount);
+
 
         var relabel = HugeLongArray.newArray(nodeCount);
         var currentCondensedRoot = nodeCount;
@@ -59,6 +61,8 @@ CondensedTree condense(ClusterHierarchy clusterHierarchy, long minClusterSize) {
         var bfsQueue = HugeLongArrayQueue.newQueue(nodeCount);
         var visited = HugeAtomicBitSet.create(clusterHierarchyRoot + 1);
 
+        size.set(currentCondensedRoot - nodeCount, nodeCount);
+
         for (var i = clusterHierarchyRoot; i >= nodeCount; i--) {
             if (visited.get(i)) {
                 continue;
@@ -88,15 +92,17 @@ CondensedTree condense(ClusterHierarchy clusterHierarchy, long minClusterSize) {
                 relabel.set(left - nodeCount, leftClusterId);
                 parent.set(leftClusterId, currentReLabel);
                 lambda.set(leftClusterId, fallingOutLambda);
+                size.set(leftClusterId - nodeCount, leftSize);
 
                 var rightClusterId = ++currentCondensedMaxClusterId;
                 relabel.set(right - nodeCount, rightClusterId);
                 parent.set(rightClusterId, currentReLabel);
                 lambda.set(rightClusterId, fallingOutLambda);
+                size.set(rightClusterId - nodeCount, rightSize);
             }
         }
 
-        return new CondensedTree(currentCondensedRoot, parent, lambda, currentCondensedMaxClusterId);
+        return new CondensedTree(currentCondensedRoot, parent, lambda, size, currentCondensedMaxClusterId, nodeCount);
     }
 
     private void fallOut(

algo/src/main/java/org/neo4j/gds/hdbscan/CondensedTree.java

@@ -27,13 +27,24 @@ class CondensedTree {
     private final long root;
     private final HugeLongArray parent;
     private final HugeDoubleArray lambda;
+    private final HugeLongArray size;
     private final long maximumClusterId;
+    private final long nodeCount;
 
-    CondensedTree(long root, HugeLongArray parent, HugeDoubleArray lambda, long maximumClusterId) {
+    CondensedTree(
+        long root,
+        HugeLongArray parent,
+        HugeDoubleArray lambda,
+        HugeLongArray size,
+        long maximumClusterId,
+        long nodeCount
+    ) {
         this.root = root;
         this.parent = parent;
         this.lambda = lambda;
+        this.size = size;
         this.maximumClusterId = maximumClusterId;
+        this.nodeCount = nodeCount;
     }
 
     long root() {
@@ -55,4 +66,8 @@ long maximumClusterId() {
     double lambda(long node) {
         return lambda.get(node);
     }
+
+    long size(long node) {
+        return size.get(node - nodeCount);
+    }
 }

algo/src/test/java/org/neo4j/gds/hdbscan/CondenseStepTest.java

@@ -43,16 +43,22 @@ void minClusterSizeTwo() {
         assertThat(condensedTree.root()).isEqualTo(7L);
         assertThat(condensedTree.maximumClusterId()).isEqualTo(11L);
 
+        assertThat(condensedTree.size(7L)).isEqualTo(7L);
 
         assertThat(condensedTree.parent(8L)).isEqualTo(7L);
         assertThat(condensedTree.lambda(8L)).isEqualTo(12d);
+        assertThat(condensedTree.size(8L)).isEqualTo(2L);
         assertThat(condensedTree.parent(9L)).isEqualTo(7L);
         assertThat(condensedTree.lambda(9L)).isEqualTo(12d);
+        assertThat(condensedTree.size(9L)).isEqualTo(5L);
 
         assertThat(condensedTree.parent(10L)).isEqualTo(9L);
+        assertThat(condensedTree.size(10L)).isEqualTo(2L);
         assertThat(condensedTree.lambda(10L)).isEqualTo(10d);
+
         assertThat(condensedTree.parent(11L)).isEqualTo(9L);
         assertThat(condensedTree.lambda(11L)).isEqualTo(10d);
+        assertThat(condensedTree.size(11L)).isEqualTo(3L);
 
         assertThat(condensedTree.fellOutOf(0L)).isEqualTo(8L);
         assertThat(condensedTree.lambda(0L)).isEqualTo(11d);
@@ -88,6 +94,8 @@ void minClusterSizeThree() {
         assertThat(condensedTree.root()).isEqualTo(7L);
         assertThat(condensedTree.maximumClusterId()).isEqualTo(7L);
 
+        assertThat(condensedTree.size(7L)).isEqualTo(7L);
+
         assertThat(condensedTree.fellOutOf(0L)).isEqualTo(7L);
         assertThat(condensedTree.lambda(0L)).isEqualTo(12d);
         assertThat(condensedTree.fellOutOf(1L)).isEqualTo(7L);