ARC 093 E - Bichrome Spanning Tree (900) - くるのプログラミング記録

900点。
天才解法というよりは理詰めで押し切れた問題なのでよかった。
リンク：E - Bichrome Spanning Tree
ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー

問題文

$N$ 頂点 $M$ 辺からなる重み付き無向グラフがあります。グラフの $i$ 番目の辺は頂点 $U_i$ と頂点 $V_i$ を結んでおり、重みは $W_i$ です。さらに、整数 $X$ が与えられます。

このグラフの各辺を白または黒に塗る方法であって、以下の条件を満たすものの個数を $10^{9} + 7$ で割ったあまりを求めてください。

白く塗られた辺と黒く塗られた辺をともに含む全域木が存在する。さらに、そのような全域木のうち最も重みが小さいものの重みは $X$ である。
ただし、全域木の重みとは、その全域木に含まれる辺の重みの和を表します。

制約

$1 \leq N \leq 1000$

$1 \leq M \leq 2000$

$1 \leq U_i, V_i \leq N (1 \leq i \leq M)$

$1 \leq W_i \leq 10^9$

$i \neq j$ のとき $(U_i, V_i) \neq (U_j, V_j)$ かつ $(U_i, V_i) \neq (V_j, U_j)$

$U_i \neq V_i (1 \leq i \leq M)$

与えられるグラフは連結である。

$1 \leq X \leq 10^{12}$

入力値はすべて整数である。

ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー
まず問題文がややこしい。サンプル見ないと何をしたい問題か分からなかった。

まず愚直を考えると全辺について白か黒かの2択があるので $2^M ( M \leq 2000)$ 通り。
ここから最小全域木を求めると (先に辺を重みでソートしておいた上で)更に $O(M)$ かかるので絶対無理。
大体 $O(M \log M + 2^MM)$ ぐらい......？(Union findを考慮してないので本当はもう少し大きそう)

辺のソートと最小全域木の構築は避けられないと思ったので、辺の塗り方を最適化する問題かなあという方針で進める。
問題文読み返してみると「白く塗られた辺と黒く塗られた辺をともに含む」とあるので白と黒をそれぞれ最低一つずつ用意する必要があるなと。

2辺決めうちして貪欲......？
これだと $M(M-1)/2$ 通り試せば良く、比較的マシになった気がするのでこの方針で進める。
この2辺をそれぞれ $A,B$ とする。

最小全域木を構築する際、どう塗り分けるにしても最小重みの辺は必ず使うことに気づいたので、2本のうち1本が決まることに(すなわち $A = \min(W_i)$ )
この最小辺 $A$ に対して、 $B$ は残りの辺について $(M-1)$ 通り試せばよく、
全体のオーダーも $O(M^2)$ となり比較的高速。おっ、解けそう。

ただ、ここからが分からなくなった。
決め打ちした2辺を必ず使う、かつ白黒に当てはめるって方針なので、例えば $A$ が黒の場合、 $B$ は白。ここで他の $B$ より軽い辺が白色でかつたくさんあった場合、決め打ちしたにも関わらず $B$ が最小全域木に含まれなくなってしまうことが起こりえる。

なので直感的には、 $B$ より軽い辺は全て $B$ と違った色であってほしい。
けど全ての $B$ より軽い辺が本当にそれで適切かといわれるとそうではないので難しい。
( $B$ より軽くても最小全域木に使わない辺に関しては白黒の2通りがあるため)

簡単な場合を考えて、色塗りを気にせず最小全域木を求めると $N-1$ 本の辺が決まる(この辺の集合を $S$ とする。)
この $S$ に属する $N-1$ 辺全ての色が同じでない場合は、残りの $M - (N-1)$ 辺(この辺の集合を $T$ とする。)をどう塗ろうとも必ず同じ最小全域木になってしまう(自明といえば自明)

なので $S$ が全て同じ色(黒あるいは白一色) or それ以外(白黒混ざっている)で場合分けができる。

・ $S$ に白黒両方の色が混ざっているとき

$(2^{N-1} - 2)(2^{(M-(N-1))})$ 通り
( 最小全域木に用いる $S$ の $N-1$ 辺に関しては(全て黒)と(全て白)を除いて $(2^{N-1} - 2)$ 通り、
　 $T$ の $M - (N-1)$ 辺に関してはそれぞれ白か黒の $(2^{(M-(N-1))})$ 通りなので)

・ $S$ が全て同じ色のとき

$T$ の $M - (N-1)$ 辺から1辺を選び、 $B$ と決め打ちする( $B = T_i ( 1 \leq i \leq M - (N - 1))$ )
$B$ と違う色と仮定し、 $A=\min(W_i)$ と $B = T_i$ を用いた最小全域木を構築する。

ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー
このときの $B( = T_i ( 1 \leq i \leq M - (N - 1)))$ の選び方だが、 $T$ の $M - (N-1)$ 辺から貪欲に軽い辺から選び、最小全域木を構築する。
使った辺は $S$ と同じ色として扱う、すなわち $T$ から除き集合 $S$ に加える。
このとき $2*(2^{(M-(N-1 + i))})$ 通り
これを繰り返し2辺決め打ち全探索をする。
ーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーーー

という嘘解法で詰まる。

$T$ から貪欲に軽い辺を $B$ として選んだからといってそれが毎回最小全域木に用いられるとは限らないので。
辺の重みの順序で $T$ から $S$ に加えていくのではなく、 $B ( = T_i ( 1 \leq i \leq M - (N - 1)))$ を用いたときの最小全域木の重みの順序で $T$ から $S$ に加えていく。そうすることで上記のような解法がまかり通る(これも気付いてみると自明な気もする)

結論としては、一番軽い辺 $A$ とそれを除いた $M - 1$ 辺から $B$ として1つを選んだ $A,B$ を用いて $M - 1$ 個の最小全域木を得る。
これらの最小全域木の重みの軽い方から処理する、このとき1度用いた $B$ は $A$ と同じ色として数え上げていく。

以下実装。

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;
using ll = long long;

#define REP(i,n) for(int i = 0; i < int(n); i++)
#define FOR(i, m, n) for(int i = int(m);i < int(n);i++)
#define ALL(obj) (obj).begin(),(obj).end()
#define fi first 
#define se second

const ll MOD = (ll)1e9 + 7;

//Class - Union Find Tree
template<class T> class Union_Find {
public:
	vector<int> prnt;
	vector<int> rank;
	vector<T> dist;

	Union_Find(int N = 1) :prnt(vector<int>(N)), rank(vector<int>(N, 0)), dist(vector<int>(N, 0)) {
		for (int i = 0; i < N; ++i) prnt[i] = i;
	}

	int root(int x) {
		if (prnt[x] == x) return x;
		else {
			int r = root(prnt[x]);
			dist[x] += dist[prnt[x]];
			return prnt[x] = r;
		}
	}

	T updated_dist(int x) {
		root(x);
		return dist[x];
	}

	bool same(int x, int y) {
		return root(x) == root(y);
	}

	bool merge(int x, int y, T d) {
		d += updated_dist(x);
		d -= updated_dist(y);
		x = root(x); y = root(y);
		if (x == y) return false;
		if (rank[x] < rank[y]) {
			swap(x, y);
			d = -d;
		}
		if (rank[x] == rank[y]) ++rank[x];
		prnt[y] = x;
		dist[y] = d;
		return true;
	}

	T diff(int x, int y) {
		return updated_dist(y) - updated_dist(x);
	}
};

int main() {
	ll N, M, X; cin >> N >> M >> X;
	vector<pair<ll, pair<int, int>>> edge(M);

	//入力
	REP(i, M) {
		int U, V;
		ll W;
		cin >> U >> V >> W;
		U--;
		V--;
		edge[i] = { W,{ U,V } };
	}

	//重みでソート
	sort(ALL(edge));

	//コーナーケース
	if (N <= 2 || M < 2) {
		cout << 0 << endl;
		return 0;
	}

	//2の累乗をメモ
	vector<ll> pow2(M + 1, 1);
	FOR(i, 1, M + 1) pow2[i] = (pow2[i - 1] * 2) % MOD;

	//M-1個の最小全域木の重み
	vector<ll> sum(M - 1, 0);
	FOR(B, 1, M) {

		//UnionFind木
		Union_Find<int> UF(N);

		//最小辺とBを用いた最小全域木
		sum[B - 1] += edge[0].fi + edge[B].fi;
		UF.merge(edge[0].se.fi, edge[0].se.se, 1);
		UF.merge(edge[B].se.fi, edge[B].se.se, 1);

		//クラスカル法による最小全域木の構築
		FOR(j, 1, M) {
			int U = edge[j].se.fi;
			int V = edge[j].se.se;
			ll W = edge[j].fi;

			//連結済ならcontinue
			if (UF.same(U, V)) continue;

			//辺をmergeし重みを加える
			UF.merge(U, V, 1);
			sum[B - 1] += W;
		}
	}

	//最小全域木の重みでソート
	sort(ALL(sum));

	//軽い方からsum[i] == Xとなるものを数え上げ
	ll ans = 0;
	REP(i, M - 1) if (sum[i] == X) (ans += ((2 * pow2[M - 2 - i]) % MOD)) %= MOD;

	cout << ans << endl;
	return 0;
}

説明がややこしくなったので実装を見るのが一番な気がする。