Code Thanks Festival 2017 H - Union Sets

問題

{1}, {2}, ..., {n}というn個の集合が与えられ、これらを併合する操作をm回行う
時刻iでは要素a[i]を持つ集合とb[i]を持つ集合を併合する
次にq個のクエリが与えられるのでそれに答える
クエリの内容: 要素xjとyjが同じ集合に併合されるのは何回目の操作後か?

解法

↓サンプル1の例
f:id:creep040:20180915185715p:plain
①要素番号を頂点、併合した時刻を辺のコストとする森をクラスカル法で作る(画像青線部分)
②頂点0を仮の根として↑の森を木にする(赤線部分)
③ダブリングを使ったLCAで2頂点間のパス上のコストの最大値(=クエリの答え)をO(logn)で求める

LCAを求めるとこまでは蟻本p292の通り
パス上のコストの最大値もそれをちょっと変えるだけで求められる
cost[k][v]: 頂点vから2^k回親を辿る途中の辺のコストの最大値
par[k][v]: 頂点vから2^k回親を辿って着く頂点の番号
とすると、cost[k+1][v]はmax(cost[k][v], cost[k][par[k][v]])で更新できる

実装

queryメソッドではLCAを求めてから頂点u,vとの深さの差だけ木を登ってる

大半がライブラリ貼っただけだけどめちゃ長い
これを機にLCAのライブラリを作った

class LCA {
private:
	static const int MAXLOGV = 20;
	static const int MAXV = 1<<MAXLOGV;
	int root;
	vector<pair<int,int> > e[MAXV];
	int par[MAXLOGV][MAXV], cost[MAXLOGV][MAXV], dep[MAXV];

public:
	void dfs(int v, int p, int d) {
		par[0][v] = p;
		dep[v] = d;
		for (int i=0; i<e[v].size(); i++) {
			if (e[v][i].first!=p) dfs(e[v][i].first, v, d+1);
			else cost[0][v] = e[v][i].second;
		}
	}
	void add(int from, int to, int dis) {
		e[from].push_back({to,dis});
		e[to].push_back({from,dis});
	}
	void init(int N, int ROOT) {
		root = ROOT;
		dfs(root, -1, 0);
		for (int k=0; k<MAXLOGV-1; k++)
			for (int v=0; v<N; v++) {
				if (par[k][v]<0) par[k+1][v] = -1, cost[k+1][v] = cost[k][v];
				else par[k+1][v] = par[k][par[k][v]], cost[k+1][v] = max(cost[k][v], cost[k][par[k][v]]);
			}
	}
	int query(int uu, int vv) {
		int u = uu, v = vv, p = -1;
		if (dep[u]>dep[v]) swap(u,v);
		for (int k=0; k<MAXLOGV; k++)
			if ((dep[v]-dep[u])>>k&1) v = par[k][v];
		if (u==v) p = u;
		else {
			for(int k=MAXLOGV-1; k>=0; k--)
				if (par[k][u]!=par[k][v]) {
					u = par[k][u];
					v = par[k][v];
				}
			p = par[0][u];
		}
		if (p==0) return -1;
		u = uu, v = vv;
		int du = dep[u] - dep[p], dv = dep[v] - dep[p], res = 0;
		for (int k=0; k<MAXLOGV; k++) {
			if (du>>k&1) chmax(res, cost[k][u]), u = par[k][u];
			if (dv>>k&1) chmax(res, cost[k][v]), v = par[k][v];
		}
		return res;
	}
};

int n, m, a, b, q;
vector<pair<int, pair<int, int> > > te;
LCA lc;

struct UF {
	vector<int> p;
	UF(int n) : p(n, -1) {};
	bool unite(int u, int v) {
		if ((u = find(u)) == (v = find(v))) return false;
		if (p[u] > p[v]) swap(u, v);
		p[u] += p[v]; p[v] = u;
		return true;
	}
	bool same(int u, int v) { return find(u) == find(v); }
	int find(int u) { return p[u] < 0 ? u : p[u] = find(p[u]); }
	int usize(int u) { return -p[find(u)]; }
};

void kruskal(int n, int m, vector<pair<int,pair<int,int> > > &e) {
	sort(e.begin(), e.end());
	UF uf(n);
	for (int i=0; i<m; ++i) {
		if (!uf.same(e[i].second.first, e[i].second.second)) {
			uf.unite(e[i].second.first, e[i].second.second);
			lc.add(e[i].second.first, e[i].second.second, e[i].first);
		}
	}
	for (int i=1; i<n; ++i)
		if (!uf.same(0,i)) uf.unite(0,i), lc.add(0,i,0);
}

signed main() {
	cin >> n >> m;
	n++;
	rep(i,m) {
		cin >> a >> b;
		te.push_back({i+1,{a,b}});
	}
	kruskal(n,m,te);
	lc.init(n,0);
	cin >> q;
	while (q--) {
		cin >> a >> b;
		cout << lc.query(a,b) << endl;
	}
}
別解

解説pdfにあった並列二分探索についてググったら感動したからこっちも書いておく

並列二分探索: 全てのクエリについての二分探索をまとめて一気にやる感じ
まずクエリごとの下限l[i]=-1と上限r[i]=mを初期化して、以下の操作をlogm回行う

priority_queueに(l+r)/2(=midとする)とクエリ番号のペアを突っ込む
→ unionfindで併合の操作を1個ずつ行い、priority_queueのtopのmidの値と今見てる時刻が同じ場合、それらが既に併合されているかチェックし、併合されていたらそのクエリの上限をmに、されていなければ下限をmにする

これでO((N+M+Q)logMlogQ)になる!
書いてから思ったけどpriority_queueは要らなくて、ただのvectorとかにすれば最後のlogQは消える
今まで使ったことないからあんま頻出ではないんだろうけどすごいなこれ🧐

実装

最初にufi(unionfind)で全部併合した状態を見て、クエリの答えが-1になるやつは弾いておいた

struct UF {
	vector<int> p;
	UF(int n) : p(n, -1) {};
	bool unite(int u, int v) {
		if ((u = find(u)) == (v = find(v))) return false;
		if (p[u] > p[v]) swap(u, v);
		p[u] += p[v]; p[v] = u;
		return true;
	}
	bool same(int u, int v) { return find(u) == find(v); }
	int find(int u) { return p[u] < 0 ? u : p[u] = find(p[u]); }
	int usize(int u) { return -p[find(u)]; }
};

int n, m, a[114514], b[114514], q, x[114514], y[114514], l[114514], r[114514];

signed main() {
	cin >> n >> m;
	UF ufi(114514);
	rep(i,m) {
		cin >> a[i] >> b[i];
		ufi.unite(a[i],b[i]);
	}
	cin >> q;
	rep(i,q) {
		cin >> x[i] >> y[i];
		if (!ufi.same(x[i],y[i])) r[i] = -1;
		else l[i] = -1, r[i] = m;
	}
	rep(t,22) {
		priority_queue<pair<int,int>, vector<pair<int,int> >, greater<pair<int,int> > > p;
		rep(i,q) if (r[i]!=-1) p.push({(l[i]+r[i])/2, i});
		UF uf(114514);
		rep(i,m) {
			uf.unite(a[i],b[i]);
			while (!p.empty() && p.top().first==i) {
				int id = p.top().second; p.pop();
				if (uf.same(x[id],y[id])) r[id] = i;
				else l[id] = i;
			}
		}
	}
	rep(i,q) cout << (r[i]==-1 ? -1 : r[i]+1) << endl;
}