When two men become friends, their sets of acquaintances merge into one set. This might remind you of the disjoint-set data structure and the union-find algorithm which can merge two disjoint sets in O(α(n)) amortized time. But this is a red herring. In fact, a more “naive” approach works much better in this problem.
Consider the standard problem solved by the union-find algorithm: we have n elements, each initially belonging to a singleton set, and we want to be able to merge two sets together and to find which set a particular element belongs to. The O(α(n)) algorithm uses union-by-rank and path compression, but there is a simpler solution.
For each element, we will remember the ID of the set it belongs to, and for each set we will remember the list of its elements. Thus, find is O(1). When we want to merge two sets, we simply take the smaller of the two, iterate over its elements, and move them to the other set one by one. This trivial algorithm actually achieves O(nlogn) total merging time. This is because we always move elements from a smaller set to a bigger set, so when an element moves, its new set is always at least twice as big as the old one. Thus, each element can move at most logn times.
Let’s use this idea in our problem. Each couple will have a pair (mi, wi), where mi is the ID of the set the man belongs to, and wi is the ID of the set the woman belongs to. So two men i and j are acquaintances if mi = mj, and two couples i and j are familiar if (mi, wi) = (mj, wj).
When a man belonging in set a and a man belonging in set b become friends, we have to merge the two sets. First, we find which of them is bigger – for example the set containing a. Then we just iterate over each couple where mi = b, and change it to a. For women, we do the same. Like above, we still move each couple at most O(logn) times.
We can do this easily if, for each a, we keep a list of all couples where mi = a, and a list of all couples where wi = a. We update these lists whenever we move any couple from one set to another. Similarly, for each a and b, we remember the number of couples with these particular numbers (a, b). Whenever this number increases, we change the total number of familiar pairs of couples appropriately.