Graphs · Medium
396. Journey to the Moon
Topics: Graph
Solve it on DeVert Campus - free code runner
Problem
There are n astronauts (0-indexed), each from exactly one country. Given p pairs, each naming two astronauts known to be from the same country (this relation is transitive - if this data implies two astronauts share a country, they do), find the number of ways to pick 2 astronauts such that they are from different countries.
Input and output
Input: The first line contains an integer n, the number of astronauts. The second line contains an integer p, the number of same-country pairs. Each of the next p lines contains two space-separated integers, astronauts known to be from the same country.
Output: Print a single integer - the number of ways to pick 2 astronauts from different countries.
Example
Input: 5 3 0 1 2 3 0 4 Output: 6
Constraints
1 <= n <= 10^5
Key observation
An astronaut who NEVER appears in any given pair is still counted as their own country of size 1 - contributing zero same-country pairs (since s(s-1)/2 = 0 for s=1), but still correctly counting toward the total n in the overall combinatorial calculation.
Intuition
Given pairs of astronauts known to share a country, find how many ways to pick 2 astronauts from DIFFERENT countries.
Brute force: Check every pair of astronauts directly, determining if they're in the same country by tracing through the given relation pairs - this costs O(n²) in the worst case, unnecessary when union-find directly groups everyone efficiently.
Optimised: Union-find groups astronauts into countries in nearly O(1) amortized time per union operation - once every country's SIZE is known, simple combinatorial arithmetic (total pairs minus same-country pairs) directly computes the answer, with no need to explicitly enumerate any pairs at all.
Brute Force - Pairwise Component Check
Build an adjacency list from the p pairs and, for every one of the n astronauts, run a BFS or DFS to discover which connected component (country) it belongs to, labeling each astronaut with a component id. Then loop over every pair of astronauts (i, j) and count the pair if their component ids differ. This works because the 'same country' relation is exactly graph reachability, so two astronauts are from different countries precisely when they lie in different connected components. It is wasteful because it explicitly examines all O(n^2) pairs even though only the component sizes actually matter.
Time: O(n^2 + p)
Space: O(n + p)
Optimal - Union-Find (DSU) with Path Compression
Use a Disjoint Set Union structure with union by rank/size and path compression to merge astronauts named in each of the p pairs; each union is answered in nearly constant amortized time (inverse-Ackermann), which is more efficient in practice than building and traversing an explicit adjacency list, especially since the pairs describe merges incrementally. After processing all pairs, read off each root's component size directly from the DSU size array, then apply the same running-total formula (total pairs minus same-country pairs) in one pass over components. This is the standard, most efficient approach for this exact 'group people into countries' style problem.
Time: O((n + p) * α(n))
Space: O(n)
Better - Component Sizes with BFS/DFS
Instead of checking every pair, find each connected component with a single BFS/DFS pass over the adjacency list built from the p pairs, recording the size of each component. Then process components one at a time, keeping a running total of astronauts seen so far; for each new component of size s, add (running total * s) to the answer before adding s to the running total. This avoids ever comparing individual pairs directly and gets the same answer using the identity that cross-country pairs equal n(n-1)/2 minus same-country pairs, computed incrementally.
Time: O(n + p)
Space: O(n + p)
Step by step
1. Use UNION-FIND: union every given same-country pair - afterward, each disjoint SET represents one country's astronauts.
2. The total number of ways to pick ANY 2 astronauts is n*(n-1)/2 - subtract the number of ways to pick 2 from the SAME country, which for a country of size s is s*(s-1)/2.
3. An astronaut who never appears in any pair is still their own country of size 1, contributing 0 same-country pairs but still counting toward n.
Complexity
Time: O(n + p * alpha(n)) - p union operations (each nearly O(1) amortized with path compression), plus a final linear pass computing each country's size.
Space: O(n), for the union-find structure.
Edge cases
an astronaut who never appears in any pair (its own country of size 1), p = 0 (every astronaut in a distinct country), and n at its maximum of 10^5 to confirm the union-find approach runs efficiently.
Interview tip
Be ready to state the two combinatorial formulas explicitly (total pairs = n(n-1)/2, same-country pairs for a group of size s = s(s-1)/2) - these are exactly what turns this problem from a graph traversal exercise into simple arithmetic once the groupings are known.