KEYENCE Programming Contest 2019 E - Connecting Cities

色々な解き方があって本当に面白い問題だと思う。
E - Connecting Cities

問題

整数N, Dと長さNの整数列A_iが与えられる。以下の完全グラフの最小全域木の重みを求めよ。

  • 頂点数はN
  • 頂点iと頂点jを結ぶコストはD|i - j| + A_i + A_j

解法

Code Festival 2017 J (Tree MST) と 距離空間上のボロノイ図 - とこはるのまとめの方針に従い、2N頂点のグラフを作り、そのグラフのヴォロノイ図を計算することで解を求める。
頂点1, ..., 2Nを持つグラフに以下のように辺を張る:

  • (i, i + N)間に重みA_i (1 <= i <= N)
  • (i + N, i + N + 1)間に重みD (1 <= i <= N - 1)

このグラフのヴォロノイ図を求めればよい。以下のようなことを考えながら実装する。

  • 頂点1, ..., Nから色1, ..., Nの液体が同じ速さで流れ始める
  • 色の違う液体は触れると触れた場所がその場で固まり、互いに混じり合わない。触れた場所以外はそのまま流れ続ける。
  • 最終的に全ての液体の動きが止まった時の液体の配置がヴォロノイ図である

これを実装するとダイクストラ法に似た実装になる。以下のデータを管理すれば良い。

  • 「色iの液体がt秒後に頂点vに到達した」というイベント。時系列順に処理するので優先度付きキューがよい。
  • 「頂点vは色iで塗られた、あるいはどの色でも塗られていない」という状態。色に順番があることにすれば、一つの頂点が2色以上で同時に塗られることはないと思って良い。

最終的なヴォロノイ図において、色と色の境界はちょうどN-1個である。それぞれの境界に対応する辺を集めたものがMSTである。
計算量は優先度付きキューを使うところが一番重く、O(N log N)である。

登場する典型

実装上の注意点

  • 今回は境界がちょうどN-1個なので、Union-Find木などを使う必要はない
  • 2色が合流するイベントを適切に処理する。色の順番を見て同じイベントを2回数えないようにするのもよいし、あえて2回数えて後で2で割るのもよい。

提出: Submission #4014722 - KEYENCE Programming Contest 2019 / キーエンス プログラミング コンテスト 2019 (Rust)

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,
        d: i64,
        a: [i64; n],
    }
    let mut g = vec![Vec::new(); 2 * n];
    for i in 0 .. n {
        g[i].push((n + i, a[i]));
        g[n + i].push((i, a[i]));
    }
    for i in 0 .. n - 1 {
        g[n + i].push((n + i + 1, d));
        g[n + i + 1].push((n + i, d));
    }
    let mut que = BinaryHeap::new();
    const INF: i64 = 1 << 50;
    let mut dp = vec![(INF, 2 * n); 2 * n];
    for i in 0 .. n {
        que.push((0, i, i));
    }
    let mut tot = 0;
    while let Some((dist, v, w)) = que.pop() {
        let dist = -dist;
        if dp[w].0 == INF {
            dp[w] = (dist, v);
            for &(u, cost) in g[w].iter() { 
                que.push((-dist - cost, v, u));
            }
        } else {
            // dp[w].1 and v meets
            if v != dp[w].1 {
                tot += dp[w].0 + dist;
            }
        }
    }
    puts!("{}\n", tot / 2);
}

まとめ

最小全域木は構造が綺麗なので、解法がたくさんあるのが面白い。(この他にも想定解法の枝刈りKruskal、Boruvka、分割統治Kruskalなどがある。)