けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

KUPC 2020 F - GRIDMST

愚直にやると  O(N^{2}) 本の辺がある問題に対して、上手にやる系の問題

問題概要

 H \times W グリッドグラフがある。

広義単調増加な正整数列  A,B,C,D があり、それぞれの長さは  H−1,W,H,W−1 となっている。

  • 頂点  (i, j) と頂点  (i+1,j) は、重みが  A_{i}+B_{j} である無向辺で結ばれている
  • 頂点  (i,j) と頂点  (i,j+1) は、重みが  C_{i}+D_{j} である無向辺で結ばれている

このグラフの最小全域木に含まれる辺の重みの和を求めよ。

制約

  •  2 \le H, W \le 10^{5}
  •  1 \le A_{i}, B_{i}, C_{i}, D_{i} \le 10^{6}

考えたこと

数列が単調増加であることから、少なくとも最上部と最左部はすべて辺で結んでしまってよいことがいえる (そうでないものも重みを悪化させることなくそうなるように入れ替えられる)。

また頂点  (i, j) に着目したとき

  •  (i-1, j) - (i, j)
  •  (i, j-1) - (i, j)

のうちの重みが小さい方と結ばれたもののみを考えればよいこともいえる。これらのうちのどちらとも結ばれてないようなやつを考えたとき、上手に入れ替えることでこれらのうちのいずれかと結ばれている状態にできる。逆に、どの  (i, j) に対しても、上記の 2 辺のうちのいずれかと結ばれていれば、全体が連結になる (帰納法で示せる)。

よって問題は、最上部と最左部を除けば、以下の問題に帰着される。


 (i, j) ( 1 \le i \le H,  1 \le j \le W に対して、

 \min(A_{i-1} + B_{j}, C_{i} + D_{j-1}

の総和を求めよ


つづき

上記の総和を素直に求めると  O(HW) の計算量となってしまう。そこで、 i を固定してあげたときの総和を高速に求めることを考える。上の式は次のように変形できる。

 \min(B_{j} - D_{j-1}, C_{i} - A_{i-1}) + A_{i-1} + D_{j-1}

ここで、 A_{i-1} + D_{j-1}] の部分の総和は簡単に計算できる。また、 M = C_{i} - A_{i-1} として、 E_{j} = B_{j+1} - D_{j} とおくと、 j = 0, 1, \dots, W-2 に対して

 \min(E_{j}, M)

の総和を計算する問題へと帰着される。これは

  •  E を小さい順にソートする
  • その中で  M 以下のものが何個あるのかは lower_bound などを用いるとわかる ( m 個とする)
  • よって、 E の累積和をあらかじめ求めておけば、以下の総和を高速に求めることができる
    •  E のうち、小さい方から  m 個の総和
    •  (W-1-m) \times M

以上を  i = 1, 2, \dots, H-1 に対してやればよい。全体の計算量は  O((H + W) \log W) となる。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int main() {
    int H, W;
    cin >> H >> W;
    vector<long long> A(H-1), B(W), C(H), D(W-1), SA(H, 0), SB(W+1, 0), SC(H+1, 0), SD(W, 0);
    for (int i = 0; i < H-1; ++i) cin >> A[i], SA[i+1] = SA[i] + A[i];
    for (int i = 0; i < W; ++i) cin >> B[i], SB[i+1] = SB[i] + B[i];
    for (int i = 0; i < H; ++i) cin >> C[i], SC[i+1] = SC[i] + C[i];
    for (int i = 0; i < W-1; ++i) cin >> D[i], SD[i+1] = SD[i] + D[i];
    vector<long long> dif(W, INF), sdif(W+1, 0);
    for (int i = 0; i < W-1; ++i) dif[i] = B[i+1] - D[i];
    sort(dif.begin(), dif.end());
    for (int i = 0; i < W; ++i) sdif[i+1] = sdif[i] + dif[i];

    long long res = 0;
    for (int i = 1; i < H; ++i) res += A[i-1] + B[0];
    for (int i = 1; i < W; ++i) res += C[0] + D[i-1];
    for (int i = 1; i < H; ++i) {
        long long lim = C[i] - A[i-1];
        long long under_num = upper_bound(dif.begin(), dif.end(), lim) - dif.begin();
        long long tmp = sdif[under_num] + lim*(W-1-under_num) + A[i-1]*(W-1) + SD[W-1];
        res += tmp;
    }
    cout << res << endl;
}