ABC255 D

問題文⌗

要するに、数列が与えられて、数列の要素に1を足したり引いたりして、すべての要素をある値($X_i$)にする問題です。今回はクエリが与えられるので、クエリ分の計算量を削減しなくてはなりません。

解法⌗

愚直に解く方法は、すべての要素をループして$X_i$との差を計算する方法です。クエリ処理に$O(Q)$、ループに$O(N)$かかるので、全体の計算量は$O(NQ)$となります。しかし、$1 \leq N,Q \leq 2 \times 10^5$なので、間に合いません。

そこで、$X_i$に対して必要な操作の最小回数を、事前に計算しておく方法を考えます。

そもそも、差とはどのように計算されるのでしょうか。
例えば、$A, B \in \mathbb{R}$の差は$|A - B|$です。$A = 3, B = 2$のとき、差は$|A - B| = 1$となります。

では、ある一つの値と、複数の値との差の合計について考えてみましょう。具体的には、数列$A = \lbrace 1,2,4,5 \rbrace$と実数$X_i = 3$との差を考えます。

先ほども言いましたが、$|1-3| + |2-3| + \ldots$と計算していては間に合いません。そこで、$A_0$の時点、$A_1$の時点$\ldots$での和を先に計算しておきます。これを累積和と言います。累積和の計算には計算量$O(N)$を要しますが、最初の一回だけ計算するので、余裕で間に合います。

数列Aの累積和をSとすると、$S = \lbrace 1, 1+2, 1+2+4, 1+2+4+5 \rbrace = \lbrace 1, 3, 7, 12 \rbrace$となります。ここで、$X_i$と、$A$のそれぞれの要素との差を一気に計算する方法を考えてみましょう。

先ほどの愚直な計算式を思い浮かべると、$|1-3| + |2-3| + |4-3| + |5-3| = 6$となります。ここで、もし絶対値が外れたらどうなるでしょうか。

$(3-1) + (3-2) + (4-3) + (5-3) = 6$ となりますね。ここで、計算量を減らすにはどうすれば良いでしょうか。

3という数が多数登場するので、分配法則でくくることができます。そうすると式は、 $$ \lbrace 3 \times 2 - (1 + 2) \rbrace + \lbrace (4 + 5) - 3 \times 2 \rbrace = 6 $$ となります。ここで、$1 + 2 = S_1, 4 + 5 = S_3 - S_2, 3=X_i$より、式は、 $$ (X_i \times 2 - S_1) + \lbrace (S_3 - S_2) - X_i \times 2 \rbrace $$ となります。つまりどういうことでしょうか。

先ほどの式は、大きく左半分と右半分に分かれます。つまり、Aにおいて$X_i$より小さい部分に関するものと、$X_i$以上の部分に関するものです。なぜ左半分と右半分に分けなければいけないのか。なぜなら、絶対値を外しているので、$X_i$より小さい$A_i$との差は$X_i - A_i$、逆に$X_i$より大きい$A_i$との差は$A_i - X_1$となるからです。

数列Aで、「どこから$A_i$が$X_i$以上になるか」を考えれば、左右別々に答えを一気に計算できるということです。

しかし、「どこから$A_i$が$X_i$以上になるか」を計算するのにも、愚直にやれば$O(N)$かかってしまいます。すなわち、全体では$O(NQ)$かかってしまい、タイムオーバー。

そこで、数列Aをソートして二分探索を行えば、$O(log{N})$で計算ができます。つまり、全体でも$O(Qlog{N})$で計算でき、余裕を持って間に合います。(厳密に言えばソートに$O(Nlog{N})$かかるので、全体では$O((N+Q)log{N})$です。)

二分探索を初めて聞いた方は検索してみてください。おそらく一番簡単なアルゴリズムです。

僕が提出したコードは以下です。

#include<bits/stdc++.h>
using namespace std;
#define rep(i, n) for(int i=0; i < (n); i++)
using ll = long long;
 
int main(void) {
    ll n, q; cin >> n >> q;
    vector<ll> a(n);
    rep(i, n) cin >> a[i];
    sort(a.begin(), a.end());
 
    vector<ll> ps(n, 0);
    rep(i, n) {
        if (i == 0) {
            ps[i] = a[i];
        }
 
        else {
            ps[i] = a[i] + ps[i - 1];
        }
    }
 
    rep(i, q) {
        ll xi; cin >> xi;
 
        ll lb = lower_bound(a.begin(), a.end(), xi) - a.begin();
 
        ll ans;
 
        if(lb >= 1) ans = xi*lb-ps[lb-1] + (ps[n-1] - ps[lb-1]) - xi * (n - 1 - lb + 1);
        else if(lb == n) ans = xi * n - ps[n-1];
        else ans = ps[n-1] - n * xi;
 
        cout << ans << endl;
    }
    return 0;
}

ABC255 D

目次

問題文⌗

解法⌗