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

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

2018 codeFlyer 本選 B - 交通費 (400 点)

ちょっと重たかったん。でも累積和には大分慣れて来たのん。lower_bound とか upper_bound がどっちだけ...ってのはまだなかなか慣れないから要修行なのん。でも今回は本当にどちらでもいいので、こういうところで思考停止できるようになるのが大事なのんな。

問題へのリンク

問題概要

 N 個の整数  X_0, X_1, \dots, X_{N-1} が与えられる。 Q 個のクエリ  c, d があって、以下のようなクエリに答えよ:

 f(x) = \min({\rm abs}(x - c), d) として、各  i に対して  f(X_i) を合算した値を答える

制約

  •  1 \le N, Q \le 10^{5}
  •  0 \le X_{i}, c, d \le 10^{9}

解法

 Q 個の各クエリごとに  N 個の  X_i を個別に計算していては  O(QN) かかってしまって間に合わないので、なんらかの高速化が必要になる。

今回の問題では

  1.  X_i \le c - d のとき、 f(X_i) = d
  2.  c - d \le X_i \le c のとき、 f(X_i) = c - X_i
  3.  c \le X_i \le c + d のとき、 f(X_i) = X_i - c
  4.  c + d \le X_i のとき、 f(X_i) = d

と場合分けできる。場合分けにおいて「=」の部分はどちらに含めても良い (が、どちらかにしないといけない)。

さて、1, 2, 3, 4 どれもまとめて計上できそうである。2 と 3 については  X の累積和が必要になるので、 X の 累積和を  S とする。

そしてそれぞれの境目も二分探索 (lower_bound とか upper_bound とか) で簡単に計算できる。lower にするか、upper にするかは、上の 1〜4 の場合分けにおいて「=」をどちらに含めるかに依る (いつも lower と upper の使い分けには頭を悩ますにだが、今回は本当にどちらでもよいので思考停止できる)。1 と 2 の境目を left, 2 と 3 の境目を mid, 3 と 4 の境目を right とすると

  1. 区間 [0, left) について  d を足しあげるので d × left
  2. 区間 [left, mid) について  c - X_i を足しあげるので c × (mid - left) - (S[mid] - S[left])
  3. 区間 [mid, right) について  X_i - c を足しあげるので (S[right] - S[mid]) - c × (right - mid)
  4. 区間 [right, N) について  d を足しあげるので d × (N - right)

あとはこれを合算する

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

long long N, Q;
vector<long long> X;
vector<long long> sum;
long long c, d;

long long solve() {
  int left = upper_bound(X.begin(), X.end(), c-d) - X.begin();
  int mid = upper_bound(X.begin(), X.end(), c) - X.begin();
  int right = upper_bound(X.begin(), X.end(), c+d) - X.begin();

  long long res14 = (left + (N - right)) * d;
  long long res2 = c*(mid-left) - (sum[mid] - sum[left]);
  long long res3 = (sum[right] - sum[mid]) - c*(right-mid);
  return res14 + res2 + res3;
}

int main() {
  while (cin >> N >> Q) {
    X.resize(N); for (int i = 0; i < N; ++i) cin >> X[i];
    sum.resize(N+1); sum[0] = 0; for (int i = 0; i < N; ++i) sum[i+1] = sum[i] + X[i];
    for (int q = 0; q < Q; ++q) {
      cin >> c >> d;
      cout << solve() << endl;
    }
  }
}