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

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

AtCoder ABC 181 E - Transformable Teacher (緑色, 500 点)

頭整理が大変だけど、考え方はわかる。

問題概要

 N は正の奇数である。 N 個の整数  H_{1}, H_{2}, \dots, H_{N} が与えられる。以下の  M 個のクエリに答えよ。

  • 各クエリでは整数  W が与えられる
  •  N+1 個の整数  H_{1}, \dots, H_{N}, W を 2 個ずつペアにして  \frac{N+1}{2} 組のペアを作る
  • それぞれのペアの「数値の差」の総和として、考えられる最小値を答えよ

制約

  •  1 \le N, M \le 2 \times 10^{5}

考えたこと

まずは 1 個のクエリに答えることを考える。これは比較的簡単にわかる。 N+1 個の整数を小さい順に並べて

 a_{1}, a_{2}, \dots, a_{N+1}

としたとき、

  •  (a_{1}, a_{2}) をペアにする (差は  a_{2} - a_{1})
  •  (a_{3}, a_{4}) をペアにする (差は  a_{4} - a_{3})
  • ...
  •  (a_{N}, a_{N+1}) をペアにする (差は  a_{N+1} - a_{N})

という風にするのが最適となる。以上で、1 つのクエリに対する問題は解けた。

クエリ処理へ

毎回のクエリにおいて  N+1 個の値のうちの  N 個は毎回共通であることに注目する。まずは、あらかじめ  H_{1}, \dots, H_{N} を小さい順に並べておこう。

このとき、クエリごとに挿入される値  W は、数列  H_{1}, \dots, H_{N} のどこかに挿入されることになる (挿入される位置は lower_bound() によって求められる)。この挿入操作による影響は、実は挿入された箇所の近辺のみに限られる。

たとえば  N = 7 で、 W が「左側に 2 個、右側に 5 個」の位置に挿入されるとしたとき

  • 「左 2 個分」は数列  H の左側から 2 個分に関する答え
  • 「右 4 個分」は数列  H の右側から 4 個分に関する答え
  • さらに  H_{2} - W

をすべて合計したものが答えになることがわかる。

これを求めるためには、次のような「左右からの累積和」を持っておくと良さそう。

  • left[  v ] :=  H の左  v 個分に関する答え ( v は偶数)
  • right[  v ] :=  H の右  v 個分に関する答え ( v は偶数)

一般に、 W が挿入されるときに、左側に  i 個あったとする。

  •  i が偶数のとき:left[  i ] + right[  N - i - 1 ] +  (H_{i} - w)
  •  i が奇数のとき:left[  i - 1 ] + right[  N - i ] +  (w - H_{i-1})

全体を通した計算量は次のようになる。

  • 数列  H のソート: O(N \log N)
  • 累積和 left, right の前処理: O(N)
  • 各クエリに答える: O(M \log N)

コード

#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 N, M;
    cin >> N >> M;
    vector<long long> X(N), W(M);
    for (int i = 0; i < N; ++i) cin >> X[i];
    for (int i = 0; i < M; ++i) cin >> W[i];
    sort(X.begin(), X.end());

    vector<long long> left(N+1, 0), right(N+1, 0);
    for (int i = 2; i < N; i += 2) {
        left[i] = left[i-2] + X[i-1] - X[i-2];
        right[i] = right[i-2] + X[N-i+1] - X[N-i];
    }
    long long res = INF;
    for (auto w : W) {
        int i = lower_bound(X.begin(), X.end(), w) - X.begin();
        if (i % 2 == 0) chmin(res, left[i] + right[N-i-1] + X[i] - w);
        else chmin(res, left[i-1] + right[N-i] + w - X[i-1]);
    }
    cout << res << endl;
}