How Big Is Your Friend Group

Medium30–35 min

Topics

GraphUnion FindDisjoint SetUnion by SizeDesign

Companies

AmazonGoogleMicrosoftUberAdobe

People start out alone, and friendships merge groups. You receive a list of operations: "merge the groups of a and b", or "how big is the group of a?". Answer every operation in order with the size of the group of a (after the merge, for merge operations).

A disjoint set (union-find) stores the group sizes at the roots of trees. Union by size hooks the smaller tree under the larger one, which keeps the trees shallow.

Test Case 1:

Input:n = 6, ops = [[1,0,1],[1,2,3],[2,3,0],[1,1,3],[2,2,0],[1,4,5],[2,4,0],[1,0,5]]

Output:[2,2,2,4,4,2,2,6]

Explanation:Merging 0–1 gives a group of 2, merging 2–3 another; asking about 3 gives 2; merging 1 and 3 makes a group of 4 (asking about 2 gives 4); 4–5 make a group of 2 (asking about 4 gives 2); the last merge 0–5 joins the groups of 4 and 2 into one of 6. As an adjacency list the friendships are 0–1, 2–3, 1–3, 4–5, 0–5.

Test Case 2:

Input:n = 3, ops = [[1,1,1],[2,2,0]]

Output:[1,1]

Explanation:Merging a person with themselves changes nothing.

Test Case 3:

Input:n = 1, ops = [[2,0,0]]

Output:[1]

Explanation:A single person is a group of size 1.

Constraints

◆1 ≤ n ≤ 10 people numbered 0 … n-1; at the start every person is alone in their own group
◆ops is a list of at most 12 operations, each a row [kind, a, b] with 0 ≤ a, b < n: kind 1 means "merge the groups of a and b" (nothing changes if they are already together), kind 2 means "ask for the size of the group of a" (b is ignored)
◆The operations are processed in the given order and every operation produces one answer
◆Return the answers: for a question the size of the group of a; for a merge the size of the group of a AFTER the merge

🚀

Try the Dry Run

Don't just read the solution — watch it execute, one step at a time.

Approach & Solutions

Brute Force — Keep Group Labels and Count the Members

Brute

Keep a group label per person (at first their own number). To merge a and b when their labels differ, sweep over all people and give everyone with b's label the label of a. To report a size, count how many people carry a's label. Both merging and counting sweep all n people, so every operation costs O(n).

TimeO(n) per operation

SpaceO(n)

Java

Try it out

class Solution {
    public int[] groupSizes(int n, int[][] ops) {
        int[] label = new int[n];
        for (int i = 0; i < n; i++) label[i] = i;
        int[] answers = new int[ops.length];
        for (int k = 0; k < ops.length; k++) {
            int a = ops[k][1], b = ops[k][2];
            if (ops[k][0] == 1 && label[a] != label[b]) {
                int from = label[b], to = label[a];
                for (int i = 0; i < n; i++) {
                    if (label[i] == from) label[i] = to;
                }
            }
            int size = 0;
            for (int i = 0; i < n; i++) {
                if (label[i] == label[a]) size++;
            }
            answers[k] = size;
        }
        return answers;
    }
}

Optimal — Disjoint Set With Union by Size and Path Compression

Optimal

Keep every group as a tree with parent[] pointers towards the root, and store the number of people in each tree at its root (size[]). find(x) walks up to the root, shortening the path as it goes (path compression). To merge, hook the root of the SMALLER group under the root of the LARGER group and add the sizes; making the big tree the root keeps trees shallow. The size of a person's group is size[find(x)], so each operation takes nearly constant time.

TimeO(α(n)) per operation