685. Redundant Connection II

/*
Using union-find with rank and path compression.

Refer: https://leetcode.com/problems/redundant-connection-ii/discuss/108045/C++Java-Union-Find-with-explanation-O(n)

Time Complexity: O(N) Space Complexity : O(N)

N = number of edges or the length of edges array.
*/
class Solution {
    class UnionFind {
        int[] parents;
        int[] ranks;

        UnionFind(int n) {
            parents = new int[n + 1];
            ranks = new int[n + 1];
        }

        public int findRoot(int x) {
            if (x != parents[x]) {
                parents[x] = findRoot(parents[x]);
            }
            return parents[x];
        }

        public boolean union(int x, int y) {
            int pX = findRoot(x);
            int pY = findRoot(y);

            if (pX == pY) {
                return false;
            }

            if (ranks[pX] >= ranks[pY]) {
                ranks[pX] = ranks[pX] == ranks[pY] ? ranks[pX] + 1 : ranks[pX];
                parents[pY] = pX;
            } else {
                parents[pX] = pY;
            }
            return true;
        }
    }

    public int[] findRedundantDirectedConnection(int[][] edges) {
        if (edges == null || edges.length == 0 || edges[0].length == 0) {
            return new int[] { -1, -1 };
        }

        int[] cand1 = new int[] { -1, -1 };
        int[] cand2 = new int[] { -1, -1 };
        UnionFind dsu = new UnionFind(edges.length);

        for (int[] e : edges) {
            int src = e[0];
            int dst = e[1];
            if (dsu.parents[dst] != 0) {
                cand1 = new int[] { dsu.parents[dst], dst };
                cand2 = new int[] { src, dst };
                e[1] = -1;
            } else {
                dsu.parents[dst] = src;
            }
        }

        // Initializing the parents array for Union-Find
        for (int i = 1; i <= edges.length; i++) {
            dsu.parents[i] = i;
        }

        for (int[] e : edges) {
            if (e[1] == -1) {
                continue;
            }
            if (!dsu.union(e[0], e[1])) {
                if (cand1[0] == -1) {
                    return e;
                } else {
                    return cand1;
                }
            }
        }

        return cand2;
    }
}