ARC093 E - Bichrome Spanning Tree

本番解けなかったけど、知識が増えた状態でやり直してみたら解けたのでよかった。
E - Bichrome Spanning Tree

問題

N頂点M辺の無向グラフがある。各辺を白か黒で塗るとき、以下の条件を満たすような辺の塗り方を求めよ。

  • 白い辺と黒い辺を1個以上含む全域木が存在する。また、そのような全域木の最小の重みはXである。
  • 1 <= N <= 1000
  • 1 <= M <= 2000
  • 1 <= (辺の重み) <= 10^9
  • 1 <= X <= 10^12
  • 与えられるグラフは連結で単純である。

解法

以降白い辺と黒い辺を1個以上含む全域木のことを白黒全域木と呼ぶことにする。
まず、問題を累積的な形に変換する。以下の値をg(y)とすると、g(X) - g(X - 1)を計算すれば良い。

  • g(y) := 重みy以下の白黒全域木が存在するような辺の塗り方の総数

これを計算するにはどうすればよいだろうか? 適当に最小全域木Sを固定した時、白黒全域木の中で最小の重みを持つものとしてあり得るのは、Sが白黒ならばSと同じ重みを持つ全域木、Sが単色ならば、Sから辺を1本取り除き1本付け加えたものに限られる。これは以下の命題を経由することで証明できる。

命題
Sを上で固定した最小全域木とする。Tを任意の全域木としたとき、S != Tならば、あるe in T \ S, f in S \ Tが存在し、T - e + fはT以下の重みを持つ。

証明
どのようなe in T \ Sに対しても、どのようにf in S \ Tをとってもw(T - e + f) > w(T)となってしまうと仮定する。このときw(e) < w(f)である。S != Tより、f' in S \ Tを一つ選ぶとあるe' in T \ Sに対してS - f' + e'は全域木である。ところが仮定よりw(e') < w(f')であるため、Sの最小性と矛盾する。(証明終)*1

この定理を使うと、任意の白黒全域木Tについて、白黒であるという条件を保ったままTから辺を消しSの辺を加える、という操作が可能になり、最終的には|T \ S| <= 1となるまでこの操作ができる。よって、Sから辺を1本取り除き1本加えた全域木だけを考えれば良い。

頂点u, vに対し、ma[u][v]を、S内のu-vパスに含まれる辺の重みの最大値とすると、辺(u, v)を加えて適切な辺を1本加える時の全域木の重みの最小値は、w(S) - ma[u][v] + (u-vの重み)であることがわかる。この時の全域木の重みがy以下となる辺全て(K)について、それらが白黒となる(つまり、白だけあるいは黒だけにならない)ことが、g(y)で数え上げられる条件と同値であることがわかる。
maは特に工夫なくO(N^2)で計算できる。集合Kの大きさ|K|はO(M)でわかる。|K|さえ分かれば、g(y) = (2^|K| - 2) * 2^{n - 1 - |K|}という形になる。(ただし、|K| = 0の時は0通り。)

実装上の注意点

  • 計算すべきものが多いので、意図通りの値が入っているか確認する。
    • maの計算をバグらせて20分溶かした

提出: Submission #3933175 - AtCoder Regular Contest 093 (Rust)

// UnionFind省略
// ModInt省略

// Stores in ma[i] the maximum weight from v to i for every i.
fn dfs_ma(mst: &[Vec<(usize, i64)>], v: usize, par: usize, cur: i64,
          ma: &mut [i64]) {
    ma[v] = cur;
    for &(w, cost) in mst[v].iter() {
        if w == par { continue; }
        dfs_ma(mst, w, v, max(cur, cost), ma);
    }
}

// Finds #ways to color g's edges s.t. there exists at least one spanning
// tree that has weight <= y.
fn calc(g: &[Vec<(usize, i64)>], ma: &[Vec<i64>], mst_weight: i64, y: i64)
        -> ModInt {
    if y < mst_weight { return ModInt::new(0); }
    let n = g.len();
    let mut relevant_edges = 0;
    let mut all_edges = 0;
    for v in 0 .. n {
        for &(w, cost) in g[v].iter() {
            if v > w { continue; }
            let diff = mst_weight - ma[v][w] + cost;
            if diff <= y { relevant_edges += 1; }
            all_edges += 1;
        }
    }
    assert!(relevant_edges >= 1);
    let two = ModInt::new(2);
    let mut ans = two.pow(relevant_edges) - two;
    ans *= two.pow(all_edges - relevant_edges);
    ans
}

fn solve() {
    let out = std::io::stdout();
    let mut out = BufWriter::new(out.lock());
    macro_rules! puts {
        ($($format:tt)*) => (write!(out,$($format)*).unwrap());
    }
    input! {
        n: usize,
        m: usize,
        x: i64,
        uvw: [(usize1, usize1, i64); m],
    }
    let mut g = vec![Vec::new(); n];
    let mut edges = Vec::new();
    for (u, v, w) in uvw {
        g[u].push((v, w));
        g[v].push((u, w));
        edges.push((w, u, v));
    }
    edges.sort();
    let mut mst = vec![Vec::new(); n];
    let mut mst_weight = 0;
    let mut uf = UnionFind::new(n);
    for (cost, u, v) in edges {
        if uf.is_same_set(u, v) { continue; }
        uf.unite(u, v);
        mst[u].push((v, cost));
        mst[v].push((u, cost));
        mst_weight += cost;
    }
    // ma[i][j]: maximum weight in the unique path i-j in mst
    let mut ma = vec![vec![-INF; n]; n];
    for i in 0 .. n {
        dfs_ma(&mst, i, n, -INF, &mut ma[i]);
    }
    puts!("{}\n", calc(&g, &ma, mst_weight, x) - calc(&g, &ma, mst_weight, x - 1));
}

まとめ

本番時のリベンジができて嬉しかったけど、それでもデバッグ含め50分かかってしまった。

*1:多分この手の議論はマトロイドで頻出の議論のはずです。参考: マトロイド - Wikipedia, 様々な全域木問題