Wcc using stack-based DFS

WeaklyConnectedComponents.java

package org.neo4j.minigds.connected_components;

import org.neo4j.minigds.api.Algorithm;
import org.neo4j.minigds.api.AlgorithmResult;
import org.neo4j.minigds.api.Graph;
import org.neo4j.minigds.utils.DirectToUndirectedGraphConverter;

import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.concurrent.atomic.AtomicInteger;

/**
 *  This algorithm calculates weakly connected components of a graph.
 */
public class WeaklyConnectedComponents implements Algorithm<ArrayList<Integer>> {

    public static final class Result implements AlgorithmResult<ArrayList<Integer>> {
        private final ArrayList<ArrayList<Integer>> components;
        private final ArrayList<Integer> nodeIdToComponentId;

        Result(ArrayList<ArrayList<Integer>> components, ArrayList<Integer> nodeIdToComponentId) {
            this.components = components;
            this.nodeIdToComponentId = nodeIdToComponentId;
        }

        @Override
        public ArrayList<Integer> getValue(int nodeId) {
            return this.components.get(this.nodeIdToComponentId.get(nodeId));
        }
    }

    static final int NOT_FOUND = -1;

    @Override
    public AlgorithmResult<ArrayList<Integer>> compute(Graph graph) {
        var undirectedGraph = DirectToUndirectedGraphConverter.convert(graph);
        int nodeCount = undirectedGraph.nodeCount();
        ArrayList<ArrayList<Integer>> components = new ArrayList<>();
        ArrayList<Integer> nodeIdToComponentId = new ArrayList<>(Collections.nCopies(nodeCount, NOT_FOUND));

        var componentId = new AtomicInteger(0);

        graph.forEachNode(nodeId -> {
            if (nodeIdToComponentId.get(nodeId) == NOT_FOUND) {
                var nextComponentId = componentId.getAndIncrement();
                var component = findConnectedComponent(undirectedGraph, nodeId, nodeIdToComponentId, nextComponentId);
                components.add(nextComponentId, component);
            }

        });

        return new Result(components, nodeIdToComponentId);
    }

    private ArrayList<Integer> findConnectedComponent(Graph g, int nodeId, ArrayList<Integer> nodeIdToComponentId, int componentId) {
        var component = new ArrayList<Integer>();

        var stack = new LinkedList<Integer>();
        stack.push(nodeId);

        while (!stack.isEmpty()) {
            var nextNode = stack.pop();
            if (nodeIdToComponentId.get(nextNode) == NOT_FOUND) {
                component.add(nextNode);
                nodeIdToComponentId.add(nextNode, componentId);
            }

            g.forEachRelationship(nextNode, (sourceId, targetId) -> {
                if (nodeIdToComponentId.get(targetId) == NOT_FOUND) {
                    stack.push(targetId);
                }
            });
        }

        return component;
    }
}