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

lower_bound() や upper_bound() を無思考で使えるとよさそう

問題へのリンク

問題概要

個の整数 が与えられる。 個のクエリ があって、以下のようなクエリに答えよ:

として、各 に対して を合算した値を答える

制約

• 
• 

解法

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

今回の問題では

1. のとき、
2. のとき、
3. のとき、
4. のとき、

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

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

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

1. 区間 [0, left) について を足しあげるので d × left
2. 区間 [left, mid) について を足しあげるので c × (mid - left) - (S[mid] - S[left])
3. 区間 [mid, right) について を足しあげるので (S[right] - S[mid]) - c × (right - mid)
4. 区間 [right, N) について を足しあげるので 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;
    }
  }
}