C++で解く「全員が友達になった最初の瞬間」問題 ― Union-Findによる効率的な解法

問題概要

あるソーシャルグループにN人の異なる人がおり、それぞれ0からN-1までの一意な整数IDが割り当てられているとします。ここにログのリストがあり、各logs[i] = [time, id_A, id_B]には、負でない整数のタイムスタンプと、2人の異なる人物のIDが含まれています。各ログは、2人が友達になった時刻を表しており、AがBと友達であれば、BもAと友達です。

人Aが人Bと「知り合い(acquainted)」であるとは、AがBと直接友達であるか、Aが「Bの知り合いである誰か」と友達であることを指します。このとき、すべての人が互いに知り合いになる最も早い時刻を求める必要があります。そのような時刻が存在しない場合は-1を返してください。

入力例と処理の流れ

入力が以下の場合を考えてみましょう。

[[20190101,0,1],[20190104,3,4],[20190107,2,3],[20190211,1,5],[20190224,2,4],[20190301,0,3],[20190312,1,2],[20190322,4,5]]
N = 6

このときの出力は 20190301 になります。理由は以下の通りです。

• タイムスタンプ 20190101: 人0と人1が友達になり、友達グループは [0,1]、[2]、[3]、[4]、[5] となります。
• タイムスタンプ 20190104: 人3と人4が友達になり、グループは [0,1]、[2]、[3,4]、[5] となります。
• タイムスタンプ 20190107: 人2と人3が友達になり、グループは [0,1]、[2,3,4]、[5] となります。
• タイムスタンプ 20190211: 人1と人5が友達になり、グループは [0,1,5]、[2,3,4] となります。
• タイムスタンプ 20190224: 人2と人4はすでに友達のため、状態は変化しません。
• タイムスタンプ 20190301: 人0と人3が友達になり、この時点で全員が互いにつながります。

解法のアプローチ(Union-Find)

この問題は、Union-Find(素集合データ構造)を使うことで効率的に解けます。手順は以下の通りです。

• find()メソッドを定義します。引数として値xを受け取り、次のように動作します。
  • parents[x]が-1であれば、xをそのまま返します。
  • parents[x] := find(parents[x]) として再帰的に根を求めます(経路圧縮)。
  • parents[x]を返します。
• メインメソッドでは、次のように処理を進めます。
  • サイズNの配列parentsとrankを定義し、parentsはすべて-1、rankはすべて1で初期化します。
  • ログをタイムスタンプ順にソートします。
  • logsの各要素iに対して以下を実行します。
    • i[1]とi[2]に対してunion(結合)を行います。
    • find(i[2])とfind(i[1])を呼び出して根を確認します。
    • 結合後のグループサイズがNになった場合(= 全員が1つのグループに統合された場合)、そのタイムスタンプi[0]を返します。
  • 最後まで条件を満たさなければ、-1を返します。

C++での実装例

以下の実装を見ると、理解がより深まるでしょう。ここでは1つの配列dsを使い、負の値で根を表すと同時に、その絶対値でグループのサイズを管理しています。

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int earliestAcq(vector<vector<int>>& logs, int N) {
        vector<int> ds (N, -1);
        sort(begin(logs), end(logs));
        for(vector<int> &k : logs) {
            if(un(k[1], k[2], ds) == N) return k[0];
        }
        return -1;
    }
    int un(int u, int v, vector<int> & ds) {
        u = find(u, ds);
        v = find(v, ds);
        if(u != v) {
            ds[v] += ds[u];
            ds[u] = v;
        }
        return -ds[v];
    }
    int find(int u, vector<int> & ds) {
        return ds[u] < 0? u : ds[u] = find(ds[u], ds);
    }
};
main(){
    vector<vector<int>> v = {
        {20190101,0,1},{20190104,3,4},{20190107,2,3},{20190211,1,5},
        {20190224,2,4},{20190301,0,3},{20190312,1,2},{20190322,4,5}
    };
    Solution ob;
    cout <<ob.earliestAcq(v, 6);
}

入力

[[20190101,0,1],[20190104,3,4],[20190107,2,3],[20190211,1,5],
[20190224,2,4],[20190301,0,3],[20190312,1,2],[20190322,4,5]]
6

出力

20190301

計算量について

ログの数をMとすると、ソートにO(M log M)、Union-Find操作全体でほぼO(M・α(N))(αはアッカーマンの逆関数で、事実上定数)の計算量となります。経路圧縮によってfindが高速化されるため、全体として非常に効率的に動作するのがポイントです。