Journey to the Moon

The problem, simply

There are N astronauts, numbered 0 to N - 1. You get a list of pairs. Each pair says "these two are from the same country". Being from the same country spreads: if 0 and 1 are, and 0 and 4 are, then 1 and 4 are too. Astronauts in no pair are alone in their own country.

The moon trip takes two astronauts from different countries. Count how many such pairs you can choose.

A story to remember it

Picture it

Examples

Way 1: the simple way

Find the countries with a search, then check every pair of astronauts.

1. Build a neighbour list from the pairs.
2. Breadth first search from every astronaut that has no country yet, and give everything it reaches the same country number.
3. For every pair i < j, count it if their country numbers differ.

import java.util.*;

class Solution {
    public long journeyToMoon(int n, int[][] pairs) {
        List<List<Integer>> adj = new ArrayList<>();
        for (int i = 0; i < n; i++) adj.add(new ArrayList<>());
        for (int[] p : pairs) {
            adj.get(p[0]).add(p[1]);
            adj.get(p[1]).add(p[0]);
        }
        int[] country = new int[n];
        Arrays.fill(country, -1);
        int id = 0;
        for (int s = 0; s < n; s++) {
            if (country[s] != -1) continue;
            Queue<Integer> line = new ArrayDeque<>();
            line.add(s);
            country[s] = id;
            while (!line.isEmpty()) {
                int u = line.poll();
                for (int v : adj.get(u)) {
                    if (country[v] == -1) { country[v] = id; line.add(v); }
                }
            }
            id++;
        }
        long count = 0;
        for (int i = 0; i < n; i++)
            for (int j = i + 1; j < n; j++)
                if (country[i] != country[j]) count++;   // check every pair
        return count;
    }
}

How long it takes: the search is O(N + P), but checking every pair is O(N^2). With 100,000 astronauts that is 5 billion checks. Extra room: the neighbour lists and country numbers, O(N + P).

Way 2: the clever way

Two upgrades.

First, build the countries with team captains (union find). Every astronaut has a parent; follow parents to the captain. For each pair, join the two captains, and keep a size for each captain.

Second, never check pairs one by one. Walk through the countries. If the new country has s astronauts and seen astronauts came before it, it adds s x seen new pairs. Then seen grows by s.

1. parent[i] = i, size[i] = 1 for everyone.
2. For each pair, find both captains. If different, hang the smaller team under the bigger one and add the sizes.
3. For each captain (an astronaut with find(i) == i): answer += size * seen, then seen += size.

import java.util.*;

class Solution {
    private int[] parent, size;

    private int find(int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];   // skip a step to keep chains short
            x = parent[x];
        }
        return x;
    }

    public long journeyToMoon(int n, int[][] pairs) {
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) { parent[i] = i; size[i] = 1; }
        for (int[] p : pairs) {
            int a = find(p[0]), b = find(p[1]);
            if (a == b) continue;
            if (size[a] < size[b]) { int t = a; a = b; b = t; }
            parent[b] = a;                   // smaller team joins the bigger one
            size[a] += size[b];
        }
        long answer = 0, seen = 0;
        for (int i = 0; i < n; i++) {
            if (find(i) != i) continue;      // only captains stand for a country
            answer += (long) size[i] * seen; // pair with everyone counted before
            seen += size[i];
        }
        return answer;
    }
}

How long it takes: each pair does two finds, which are almost constant time, and the counting loop is O(N). Total about O(N + P). Extra room: the parent and size arrays, O(N).

Watch out for

• Forgetting astronauts in no pair. Each one is a country of size 1, and the counting loop includes them because they are their own captain.
• Adding the same pair twice. The running total counts each cross-country pair exactly once, when the second of the two countries is reached.
• Using size[i] of a non-captain. Only the captain's size is kept up to date.
• Another way to count: all pairs minus same-country pairs, N(N-1)/2 - sum of s(s-1)/2. Same answer, same speed.

What interviewers ask next

• Can you do it with depth first search instead of union find? Yes. Search every unvisited astronaut, count how many it reaches, and use the same running total.
• Pick three astronauts all from different countries? Keep two running totals: pairs seen so far and astronauts seen so far. Each country of size s adds s x pairsSoFar triples.
• Pairs arrive one at a time and you must answer after each? Union find shines here. When two countries of sizes a and b merge, the answer drops by exactly a x b.
• Why union by size? It keeps every chain of parents short, about log N at worst, even without path compression.

Remember this

Grown-up words: astronauts are vertices, same-country pairs are edges, countries are connected components, team captains are union find roots, hanging the smaller team under the bigger is union by size, and the running total is a prefix sum trick for counting cross pairs.