Merge pull request #6870 from IoannisPanagiotas/spanning-tree-bugs-23

[2.3] K-Spanning tree bugs

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/spanningtree/SpanningTree.java

@@ -20,7 +20,7 @@
 package org.neo4j.gds.spanningtree;
 
 import org.apache.commons.lang3.builder.EqualsBuilder;
-import org.neo4j.gds.api.RelationshipConsumer;
+import org.neo4j.gds.api.RelationshipWithPropertyConsumer;
 import org.neo4j.gds.core.utils.paged.HugeDoubleArray;
 import org.neo4j.gds.core.utils.paged.HugeLongArray;
 
@@ -62,21 +62,24 @@ public double totalWeight() {
         return totalWeight;
     }
 
-    public HugeLongArray parentArray() {return parent;}
+    public HugeLongArray parentArray() {
+        return parent;
+    }
 
     public long parent(long nodeId) {return parent.get(nodeId);}
 
     public double costToParent(long nodeId) {
         return costToParent.get(nodeId);
     }
 
-    public void forEach(RelationshipConsumer consumer) {
+    public void forEach(RelationshipWithPropertyConsumer consumer) {
         for (int i = 0; i < nodeCount; i++) {
-            final long parent = this.parent.get(i);
+            long parent = this.parent.get(i);
+            double cost = this.costToParent(i);
             if (parent == -1) {
                 continue;
             }
-            if (!consumer.accept(parent, i)) {
+            if (!consumer.accept(parent, i, cost)) {
                 return;
             }
         }

alpha/alpha-algo/src/main/java/org/neo4j/gds/impl/spanningtree/KSpanningTree.java

@@ -19,8 +19,11 @@
  */
 package org.neo4j.gds.impl.spanningtree;
 
+import com.carrotsearch.hppc.BitSet;
+import org.jetbrains.annotations.NotNull;
 import org.neo4j.gds.Algorithm;
 import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.core.utils.paged.HugeDoubleArray;
 import org.neo4j.gds.core.utils.paged.HugeLongArray;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.queue.HugeLongPriorityQueue;
@@ -45,8 +48,6 @@ public class KSpanningTree extends Algorithm<SpanningTree> {
     private final long startNodeId;
     private final long k;
 
-    private SpanningTree spanningTree;
-
     public KSpanningTree(
         Graph graph,
         DoubleUnaryOperator minMax,
@@ -71,38 +72,248 @@ public SpanningTree compute() {
             startNodeId,
             progressTracker
         );
+
         prim.setTerminationFlag(getTerminationFlag());
         SpanningTree spanningTree = prim.compute();
-        HugeLongArray parent = spanningTree.parentArray();
-        long parentSize = parent.size();
-        HugeLongPriorityQueue priorityQueue = minMax == Prim.MAX_OPERATOR ? HugeLongPriorityQueue.min(parentSize) : HugeLongPriorityQueue.max(
-            parentSize);
-        progressTracker.beginSubTask(parentSize);
-        for (long i = 0; i < parentSize && terminationFlag.running(); i++) {
-            long p = parent.get(i);
-            if (p == -1) {
-                continue;
-            }
-            priorityQueue.add(i, graph.relationshipProperty(p, i, 0.0D));
-            progressTracker.logProgress();
-        }
+
+        var outputTree = growApproach(spanningTree);
         progressTracker.endSubTask();
-        progressTracker.beginSubTask(k - 1);
-        // remove k-1 relationships
-        for (long i = 0; i < k - 1 && terminationFlag.running(); i++) {
-            long cutNode = priorityQueue.pop();
-            parent.set(cutNode, -1);
-            progressTracker.logProgress();
+        return outputTree;
+    }
+
+    @NotNull
+    private HugeLongPriorityQueue createPriorityQueue(long parentSize, boolean pruning) {
+        boolean minQueue = minMax == Prim.MIN_OPERATOR;
+        //if pruning, we remove the worst (max if it's a minimization problem)
+        //therefore we flip the priority queue
+        if (pruning) {
+            minQueue = !minQueue;
         }
-        progressTracker.endSubTask();
-        this.spanningTree = prim.getSpanningTree();
-        progressTracker.endSubTask();
-        return this.spanningTree;
+        HugeLongPriorityQueue priorityQueue = minQueue
+            ? HugeLongPriorityQueue.min(parentSize)
+            : HugeLongPriorityQueue.max(parentSize);
+        return priorityQueue;
     }
 
     @Override
     public void release() {
         graph = null;
-        spanningTree = null;
     }
+
+    private double init(HugeLongArray parent, HugeDoubleArray costToParent, SpanningTree spanningTree) {
+        graph.forEachNode((nodeId) -> {
+            parent.set(nodeId, spanningTree.parent(nodeId));
+            costToParent.set(nodeId, spanningTree.costToParent(nodeId));
+            return true;
+        });
+        return spanningTree.totalWeight();
+    }
+
+
+    private SpanningTree growApproach(SpanningTree spanningTree) {
+
+        //this approach grows gradually the MST found in the previous step
+        //when it is about to get larger than K, we crop the current worst leaf if the new value to be added
+        // is actually better
+        if (spanningTree.effectiveNodeCount() < k)
+            return spanningTree;
+
+        HugeLongArray outDegree = HugeLongArray.newArray(graph.nodeCount());
+
+        HugeLongArray parent = HugeLongArray.newArray(graph.nodeCount());
+        HugeDoubleArray costToParent = HugeDoubleArray.newArray(graph.nodeCount());
+
+        init(parent, costToParent, spanningTree);
+        double totalCost = 0;
+        var priorityQueue = createPriorityQueue(graph.nodeCount(), false);
+        var toTrim = createPriorityQueue(graph.nodeCount(), true);
+
+        //priority-queue does not have a remove method
+        // so we need something to know if a node is still a leaf or not
+        BitSet exterior = new BitSet(graph.nodeCount());
+        //at any point, the tree has a root we mark its neighbors in this bitset to avoid looping to find them
+        BitSet rootNodeAdjacent = new BitSet(graph.nodeCount());
+        //we just save which nodes are in the final output and not (just to do clean-up; probably can be avoided)
+        BitSet included = new BitSet(graph.nodeCount());
+
+        priorityQueue.add(startNodeId, 0);
+        long root = startNodeId; //current root is startNodeId
+        long nodesInTree = 0;
+        progressTracker.beginSubTask(graph.nodeCount());
+        while (!priorityQueue.isEmpty()) {
+            long node = priorityQueue.top();
+            progressTracker.logProgress();
+            double associatedCost = priorityQueue.cost(node);
+            priorityQueue.pop();
+            long nodeParent = parent.get(node);
+
+            boolean nodeAdded = false;
+            if (nodesInTree < k) { //if we are smaller, we can just add it no problemo
+                nodesInTree++;
+                nodeAdded = true;
+            } else {
+                var nodeToTrim = findNextValidLeaf(toTrim, exterior); //a leaf node with currently theworst cost
+                if (parent.get(node) == nodeToTrim) {
+                    //we cannot add it, if we're supposed to remove its parent
+                    //TODO: should be totally feasible to consider the 2nd worst then.
+                    continue;
+                }
+
+                boolean shouldMove = moveMakesSense(associatedCost, toTrim.cost(nodeToTrim), minMax);
+
+                if (shouldMove) {
+                    nodeAdded = true;
+
+                    double value = toTrim.cost(nodeToTrim);
+                    toTrim.pop();
+
+                    long parentOfTrimmed = parent.get(nodeToTrim);
+                    included.clear(nodeToTrim); //nodeToTrim is removed from the answer
+                    clearNode(nodeToTrim, parent, costToParent);
+                    totalCost -= value; //as well as its cost from the solution
+
+                    if (root != nodeToTrim) { //we are not removing the actual root
+                        //reduce degree of parent
+                        outDegree.set(parentOfTrimmed, outDegree.get(parentOfTrimmed) - 1);
+                        long affectedNode = -1;
+                        double affectedCost = -1;
+                        long parentOutDegree = outDegree.get(parentOfTrimmed);
+                        if (parentOfTrimmed == root) { //if its parent is the root
+                            rootNodeAdjacent.clear(nodeToTrim); //remove the trimmed child
+                            if (parentOutDegree == 1) { //root becomes a leaf
+                                assert rootNodeAdjacent.cardinality() == 1;
+                                //get the single sole child of root
+                                var rootChild = rootNodeAdjacent.nextSetBit(0);
+                                affectedNode = root;
+                                affectedCost = costToParent.get(rootChild);
+                            }
+                        } else {
+                            if (parentOutDegree == 0) { //if parent becomes a leaf
+                                affectedNode = parentOfTrimmed;
+                                affectedCost = costToParent.get(parentOfTrimmed);
+                            }
+                        }
+                        if (affectedNode != -1) { //if a node has been converted to a leaf
+                            updateExterior(affectedNode, affectedCost, toTrim, exterior);
+                        }
+                    } else {
+                        //the root is removed, long live the new root!
+                        assert rootNodeAdjacent.cardinality() == 1;
+                        //the new root is the single sole child of old root
+                        var newRoot = rootNodeAdjacent.nextSetBit(0);
+                        rootNodeAdjacent.clear(); //empty everything
+                        //find the children of the new root (this can happen once per node)
+
+                        fillChildren(newRoot, rootNodeAdjacent, parent, included);
+
+                        root = newRoot;
+                        //set it as root
+                        clearNode(root, parent, costToParent);
+                        //check if root is a degree-1 to add to exterior
+                        if (outDegree.get(root) == 1) {
+                            //get single child
+                            var rootChild = rootNodeAdjacent.nextSetBit(0);
+                            priorityQueue.add(root, costToParent.get(rootChild));
+                            exterior.set(root);
+                        }
+                    }
+                }
+            }
+            if (nodeAdded) {
+                included.set(node); // include it in the solution (for now!)
+                totalCost += associatedCost; //add its associated cost to the weight of tree
+                if (nodeParent == root) { //if it's parent is the root, update the bitset
+                    rootNodeAdjacent.set(node);
+                }
+                if (node != root) { //this only happens for startNode to be fair
+                    //the node's parent gets an update in degree
+                    outDegree.set(nodeParent, outDegree.get(nodeParent) + 1);
+                    exterior.clear(nodeParent); //and remoed from exterior if included
+                }
+                //then the node  (being a leaf) is added to the trimming priority queu
+                toTrim.add(node, associatedCost);
+                exterior.set(node); //and the exterior
+                relaxNode(node, priorityQueue, parent, spanningTree);
+
+            } else {
+                clearNode(node, parent, costToParent);
+            }
+        }
+        //post-processing step: anything not touched is reset to -1
+        pruneUntouchedNodes(parent, costToParent, included);
+        progressTracker.endSubTask();
+        return new SpanningTree(root, graph.nodeCount(), k, parent, costToParent, totalCost);
+
+    }
+
+    private void pruneUntouchedNodes(HugeLongArray parent, HugeDoubleArray costToParent, BitSet included) {
+        graph.forEachNode(nodeId -> {
+            if (!included.get(nodeId)) {
+                clearNode(nodeId, parent, costToParent);
+            }
+            return true;
+        });
+    }
+
+    private void clearNode(long node, HugeLongArray parent, HugeDoubleArray costToParent) {
+        parent.set(node, -1);
+        costToParent.set(node, -1);
+    }
+
+    private boolean moveMakesSense(double cost1, double cost2, DoubleUnaryOperator minMax) {
+        if (minMax == Prim.MAX_OPERATOR) {
+            return cost1 > cost2;
+        } else {
+            return cost1 < cost2;
+        }
+    }
+
+    private void updateExterior(long affectedNode, double affectedCost, HugeLongPriorityQueue toTrim, BitSet exterior) {
+        if (!toTrim.containsElement(affectedNode)) {
+            toTrim.add(affectedNode, affectedCost); //add it to pq
+        } else {
+            //it is still in the queue, but it is not a leaf anymore, so it's value is obsolete
+            toTrim.set(affectedNode, affectedCost);
+        }
+        exterior.set(affectedNode); //and mark it in the exterior
+    }
+
+    private long findNextValidLeaf(HugeLongPriorityQueue toTrim, BitSet exterior) {
+        while (!exterior.get(toTrim.top())) { //not valid frontier nodes anymore, just ignore
+            toTrim.pop(); //as we said, pq does not have a direct remove method
+        }
+        return toTrim.top();
+    }
+
+    private void fillChildren(long newRoot, BitSet rootNodeAdjacent, HugeLongArray parent, BitSet included) {
+        graph.forEachRelationship(newRoot, (s, t) -> {
+            //relevant are only those nodes which are currently
+            //in the k-tree
+            if (parent.get(t) == s && included.get(t)) {
+                rootNodeAdjacent.set(t);
+            }
+            return true;
+        });
+    }
+
+    private void relaxNode(
+        long node,
+        HugeLongPriorityQueue priorityQueue,
+        HugeLongArray parent,
+        SpanningTree spanningTree
+    ) {
+        graph.forEachRelationship(node, (s, t) -> {
+            if (parent.get(t) == s) {
+                //TODO: work's only on mst edges for now (should be doable to re-find an k-MST from whole graph)
+                if (!priorityQueue.containsElement(t)) {
+                    priorityQueue.add(t, spanningTree.costToParent(t));
+                }
+
+            }
+            return true;
+        });
+    }
+
 }
+
+

alpha/alpha-algo/src/main/java/org/neo4j/gds/impl/spanningtree/KSpanningTreeAlgorithmFactory.java

@@ -52,8 +52,7 @@ public Task progressTask(
             ) {
             return Tasks.task(
                 taskName(),
-                Tasks.leaf("SpanningTree", graph.nodeCount()),
-                Tasks.leaf("Add relationship weights"),
+                Tasks.leaf("SpanningTree", graph.relationshipCount()),
                 Tasks.leaf("Remove relationships")
             );
         }