本番見すらしなかった問題だけど、解説を見たりドワコン5th本戦C問題の準備をしたりして理解が深まったので解けた。
J - Tree MST
問題
辺に重みがついているN頂点の木、および長さNの数列Xが与えられる。このとき、以下の完全グラフの最小全域木の重みを求めよ。
- 頂点数はN
- (u, v)の重みは、X[u] + X[v] + (uv間の距離)である。
制約
- 2 <= N <= 2 * 10^5
- 1 <= X[i] <= 10^9
- 重みは1から10^9までの整数
- 入力は全て整数
解法
最小全域木なので、PrimかKruskalかブルーフカ法のどれかをまず検討すべきである。
今回はブルーフカ法を使う。頂点や連結成分ごとに一番近い点を求めることが、全て合わせてO(N)時間程度で求めることができるためである。
連結成分を色だと思うと、各頂点uについてX[v] + (uv間の距離)が最小となる点vとその色、および2番手の色について同じものが分かれば良い。これは全方位木DPでできる。
それができたら、連結成分ごとに一番近い別の連結成分を求めて、それらに対し辺を張る。
計算量はO(N log N)である。ただしVecを使いまくるので、定数倍は重い。
実装上の注意点
- オーバーフローに気をつける。今回は全ての木の辺をつなぐという実行可能な解が存在し、それのコストは高々2N * max(X_i) + \sum (辺の重み) <= 10^15なので、辺の重みをi64で管理すればオーバーフローはしない。
- 同じ重みを持つ辺のタイブレイクは一貫させる!(そうでないと、例えば辺が全て同じ重みの3角形で壊れる)
- 頂点番号でよい
- ブルーフカ法の各ステップでは、各頂点からではなく、各連結成分から一番近い連結成分に向かって辺を生やす!!
今の所各頂点からで失敗する簡単な例はわからないが、その方法で実装して提出してみるとWAになった- 追記: 例は下にある。
- mergeはO(N)でできるが、サボってsortを使うとO(N log N)になる。結構リスキーなのでO(N)にしたいところ。
提出: Submission #3917559 - CODE FESTIVAL 2017 Final (Rust)
type Cost = (i64, usize, usize); fn merge(pool: Vec<(Cost, usize)>) -> Vec<(Cost, usize)> { if pool.len() == 0 { return pool; } let &mi = pool.iter().min().unwrap(); let mut ans = vec![mi]; let mi2 = pool.iter().filter(|&&(_, color)| color != mi.1).min(); match mi2 { None => (), Some(&x) => ans.push(x), } ans } fn dfs1(v: usize, par: usize, x: &[i64], g: &[Vec<(usize, i64)>], color: &[usize], dp1: &mut [Vec<(Cost, usize)>]) { let mut pool = Vec::new(); for &(w, cost) in g[v].iter() { if w == par { continue; } dfs1(w, v, x, g, color, dp1); for &((a, _, u), b) in dp1[w].iter() { pool.push(((a + cost, v, u), b)); } pool.push(((cost + x[w], v, w), color[w])); } dp1[v] = merge(pool); } fn dfs2(v: usize, par: usize, x: &[i64], g: &[Vec<(usize, i64)>], color: &[usize], dp1: &[Vec<(Cost, usize)>], dp2: &mut [Vec<(Cost, usize)>], pardist: i64) { let mut pool = Vec::new(); let n = g.len(); if par < n { for &((d, _, dec), c) in dp2[par].iter() { pool.push(((d + pardist, v, dec), c)); } pool.push(((pardist + x[par], v, par), color[par])); } pool.append(&mut dp1[v].clone()); dp2[v] = merge(pool); for &(w, cost) in g[v].iter() { if w == par { continue; } dfs2(w, v, x, g, color, dp1, dp2, cost); } } const INF: i64 = 1 << 60; fn calc(x: &[i64], g: &[Vec<(usize, i64)>]) -> i64 { let n = x.len(); let mut uf = UnionFind::new(n); let mut conn = n; let mut tot: i64 = 0; while conn > 1 { let oldconn = conn; let mut color = vec![0; n]; for i in 0 .. n { color[i] = uf.root(i); } let mut dp1 = vec![Vec::new(); n]; let mut dp2 = vec![Vec::new(); n]; dfs1(0, n, &x, &g, &color, &mut dp1); dfs2(0, n, &x, &g, &color, &dp1, &mut dp2, 0); let mut edges = vec![(INF, n, n); n]; for i in 0 .. n { let ((cost, anc, dec), _idx) = if !uf.is_same_set((dp2[i][0].0).1, (dp2[i][0].0).2) { dp2[i][0] } else { dp2[i][1] }; let cost = cost + x[i]; let (anc, dec) = if anc > dec { (dec, anc) } else { (anc, dec) }; assert_ne!(anc, dec); let re = &mut edges[uf.root(i)]; *re = min(*re, (cost, anc, dec)); } for (cost, anc, dec) in edges { if cost < INF && !uf.is_same_set(anc, dec) { tot = tot.checked_add(cost).unwrap(); uf.unite(anc, dec); conn -= 1; } } assert_ne!(oldconn, conn); } tot }
まとめ
連結成分ごとではなく頂点ごとに一番近い点へ辺を伸ばすのも同じくらいうまくいきそうなのに、なぜダメなのだろうか? (簡単な反駁用ケースが欲しい)
追記 (2019/01/03 19:08): 以下のようなケースで反駁できる。ここで0と1は重み1の辺で連結しており、他の頂点はすべて1点からなる連結成分をなしている。仮にここで各頂点から最も近い連結成分への辺を追加すると、重み3, 2, 1000の辺が追加されるが、明らかにこれは最小全域木にはならない。
