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*(*n*log*n*) 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 log*n* times.

Let’s use this idea in our problem. Each couple will have a pair (*m*_{i}, *w*_{i}), where *m*_{i} is the ID of the set the man belongs to, and *w*_{i} is the ID of the set the woman belongs to. So two men *i* and *j* are acquaintances if *m*_{i} = *m*_{j}, and two couples *i* and *j* are familiar if (*m*_{i}, *w*_{i}) = (*m*_{j}, *w*_{j}).

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 *m*_{i} = *b*, and change it to *a*. For women, we do the same. Like above, we still move each couple at most *O*(log*n*) times.

We can do this easily if, for each *a*, we keep a list of all couples where *m*_{i} = *a*, and a list of all couples where *w*_{i} = *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.