Hackerrank's Merging Communities

Question

Challenge can be found here

People connect with each other in a social network. A connection between Person I and Person J is represented as M I J. When two persons belonging to different communities connect, the net effect is the merger of both communities which I and J belongs to.

At the beginning, there are N people representing N communities. Suppose person 1 and 2 connected and later 2 and 3 connected, then 1,2, and 3 will belong to the same community.

There are two type of queries:

M I J => communities containing person I and J merged (if they belong to different communities).

Q I => print the size of the community to which person belongs.

Input Format

The first line of input will contain integers N and Q, i.e. the number of people and the number of queries. The next Q lines will contain the queries.

Constraints :

1 <= N <= 100000

1 <= Q <= 200000

My code times out for 6 / 9 test-cases, so obviously this can be improved a lot.

Any suggestions would be much appreciated.

import java.io.*;
import java.util.*;

class Solution {
    public static void main(String[] args) {
        Scanner s = new Scanner(System.in);
        int N = s.nextInt();
        int Q = s.nextInt();
        DisjointSet ds = new DisjointSet();

        for(int i = 1; i <= N; i++) {
            ds.makeSet(i);
        }

        for(int i = 0; i < Q; i++) {
            int command = s.next().charAt(0);
            if(command == 'Q') {
                System.out.println(ds.getSetSize(s.nextInt()));
            } else {
                ds.union(s.nextInt(), s.nextInt());
            }
        }
    }
}


class DisjointSet {
    private List<Map<Integer, Set<Integer>>> disjointSet;

    DisjointSet() {
        disjointSet = new ArrayList<Map<Integer, Set<Integer>>>();
    }

    public void makeSet(int element) {
        Map<Integer, Set<Integer>> map = new HashMap<>();
        Set<Integer> set = new HashSet<>();
        set.add(element);
        map.put(element, set);
        disjointSet.add(map);
    }

    public void union(int a, int b) {
        int first = find(a);
        int second = find(b);
        Set<Integer> firstSet = null;
        Set<Integer> secondSet = null;

        for(int i = 0; i < disjointSet.size(); i++) {
            Map<Integer, Set<Integer>> map = disjointSet.get(i);
            if(map.containsKey(first)) {
                firstSet = map.get(first);
            } else if(map.containsKey(second)) {
                secondSet = map.get(second);
            }
        }

        if(firstSet != null && secondSet != null) {
            firstSet.addAll(secondSet);
        }

        for(int i = 0; i < disjointSet.size(); i++) {
            Map<Integer, Set<Integer>> map = disjointSet.get(i);
            if (map.containsKey(first)) {
                map.put(first, firstSet);
            } else if (map.containsKey(second)) {
                map.remove(second);
                disjointSet.remove(i);
            }
        }
    }

    public int find(int n) {
        for(int i = 0; i < disjointSet.size(); i++) {
            Map<Integer, Set<Integer>> map = disjointSet.get(i);
            Set<Integer> keySet = map.keySet();
            for(Integer key : keySet) {
                Set<Integer> set = map.get(key);
                if(set.contains(n)) {
                    return key;
                }
            }
        }
        return -1;
    }

    public int getSetSize(int n) {
        for(int i = 0; i < disjointSet.size(); i++) {
            Map<Integer, Set<Integer>> map = disjointSet.get(i);
            Set<Integer> keySet = map.keySet();
            for(Integer key : keySet) {
                Set<Integer> set = map.get(key);
                if(set.contains(n)) {
                    return set.size();
                }
            }
        }
        return -1;
    }
}

Answer 1

Simplifying

        for(int i = 0; i < disjointSet.size(); i++) {
            Map<Integer, Set<Integer>> map = disjointSet.get(i);
            if(map.containsKey(first)) {
                firstSet = map.get(first);
            } else if(map.containsKey(second)) {
                secondSet = map.get(second);
            }
        }

Consider instead

        for (Map<Integer, Set<Integer>> map : disjointSet) {
            if (firstSet == null) {
                firstSet = map.get(first);
            }

            if (secondSet == null) {
                secondSet = map.get(second);
            }

            if (firstSet != null && secondSet != null) {
                break;
            }
        }

You don't need i. Java can iterate over a collection directly, which saves us a line of code.

You don't need to check containsKey before doing get. If containsKey is false, then get will return null.

You might as well stop as soon as both sets are found. You don't need to keep going.

A linear scan over all the values is something like \$\mathcal{O}(N \cdot Q)\$, assuming the number of merges is proportional to the number of queries. Because you have to do one linear scan for each merge. However, there is an alternative solution that is closer to \$\mathcal{O}(Q)\$.

Alternative data structure

        DisjointSet ds = new DisjointSet();

Consider

        DisjointSet ds = new DisjointSet(N + 1);

Now you can change from

    private List<Map<Integer, Set<Integer>>> disjointSet;

to something fixed size. The normal representation for a disjoint set is a parent pointer tree with ranks added. In this case, we need to track the sizes and can use those to represent the rank. Consider

public class DisjointSet {

    private int[] parents;
    private int[] sizes;

    DisjointSet(int N) {
        parents = new int[N];
        sizes = new int[N];
    }

    public void makeSet(int i) {
        parents[i] = i;
        sizes[i] = 1;
    }

    public void union(int a, int b) {
        int first = find(a);
        int second = find(b);

        // if already part of the same set, no need to union
        if (first == second) {
            return;
        }

        if (sizes[first] < sizes[second]) {
            parents[first] = second;
            sizes[second] += sizes[first];
        } else {
            parents[second] = first;
            sizes[first] += sizes[second];
        }
    }

    public int find(int i) {
        // if not the root
        if (parents[i] != i) {
            // Make the parent the root, so that it will recurse at most once
            // on subsequent calls.
            parents[i] = find(parents[i]);
        }

        // Return the root (as the parent is always the root by this point).
        return parents[i];
    }

    public int getSetSize(int i) {
        return sizes[find(parents[i])];
    }

}

The parents array stores the immediate parent of any element. If an element is alone in a set, it is its own parent. Root nodes (terminators) are also their own parents.

For a root node, the sizes array holds the size of the set. For other nodes, it holds junk data.

This uses both union-by-rank and path compression for an optimal runtime. It's also reasonably efficient in space, requiring only two integers for each element.