Mister雑記

競プロやります。

ARC100-C 「Linear Approximation」(300)

離散的ながらこれも立派な微分

問題リンク

概要

長さ Nの数列 Aが与えられる。 bに対して「悲しさ」 s(b)

 \displaystyle s(b) = \sum_{i=1}^N \left| A_i - (b + i) \right|

によって定めたとき、悲しさの最小値を求めよ。

制約

  •  1 \leq N \leq 2 \times 10^5
  •  1 \leq A_i \leq 10^9

解説

 s(b) A_iのindexに依存しているが、上の定義式は書き直すと

 \displaystyle s(b) = \sum_{i=1}^N \left| (A_i - i) - b \right|

となるため、 B_i = A_i - iなる数列 Bを新たに作れば

 \displaystyle s(b) = \sum_{i=1}^N \left| B_i - b \right|

となり、indexに依存しない形になる。これでソートなどの順番を変える操作も行えるようになるので、以降 Bソート済みとする。

あとはこの「各要素の bからの距離」の総和を最小化したいわけだが、ここで Bの各値を二次元平面にプロットしてみる。横軸を i、縦軸を B_iとすると、以下のような具合になる。

さて、 bが上のように B_k B_{k+1}の間にあるとき、 bを1増やすと s(b)はどれだけ変化するだろうか?

結果は上のように、

  •  i \in [1, k]については距離が1増加
  •  i \in [k + 1, N]については距離が1減少

となるため、 s(b+1) - s(b) = k - (N - k)となる。ここで d(k) = 2k - Nとしておこう。

 k \dfrac{N}{2}前後で場合分けすると以下のようになる。

 k \lt \dfrac{N}{2}  k \gt \dfrac{N}{2}
 d(k) \lt 0  d(k) \gt 0
 s(b+1) \lt s(b)  s(b+1) \gt s(b)
 bを増やした方が得  bを減らした方が得

ここから k \fallingdotseq \dfrac{N}{2} s(b)が最小になることが分かるが、厳密にはどうなるだろうか。これは Nが偶数のケースと奇数のケースで場合分けする必要があるが、結果だけ言うとどちらの場合も

 \displaystyle b = B_{\frac{N + 1}{2}}

として問題ない(小数切り捨て)。0-indexedの場合はb = B[N / 2]となる。

あとはこのときの s(b)を実際に計算すればACとなる。

計算量はソートがボトルネックとなり O(N \log N)。中央値を O(N)で出すアルゴリズムもあるらしいが、そこまでする必要はないだろう。

実装例

#include <algorithm>
#include <iostream>

using namespace std;
using ll = long long;

int main() {
    ll N;
    cin >> N;
    ll A[N];
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        A[i] -= i + 1;
    }

    sort(A, A + N);
    ll b = A[N / 2];

    ll ans = 0;
    for (int i = 0; i < N; ++i) {
        ans += abs(A[i] - b);
    }

    cout << ans << endl;
    return 0;
}

制約から、intではなくlong longを使わないとオーバーフローするので注意。

別解

なおこの問題、他にも別解があるのでそちらも軽く紹介しておく。

その1: 三分探索

先に上げた表を基準に s(b)のグラフを書いてみると、これは下に凸のグラフになっていることが分かる。こういった「凸性があるグラフについて極値を求めたい」というときに使えるのが「三分探索」というアルゴリズムである。ここでは詳しく解説しないが、興味がある方は調べてみるといいだろう。

このアルゴリズムを使う問題を1つ載せておく。ネタバレになるので問題名は伏せておく。

ちなみにこの解法でも探索で使う値を求めるのに O(N)、探索は二分探索同様 O(\log N)かかるので全体計算量は O(N \log N)となり、上の解法と同じになる。

その2: 累積和

「最適な b B_1 \sim B_Nのどれかだろう」と目星をつければ、累積和を利用することで全探索が可能。具体的には b = B_kの場合

 \displaystyle s(B_k) = ((sum_N - sum_k) - (N - k)B_k) + ((k - 1)B_k - sum_{k - 1})

となる。なお sum_i = B_1 + \dots + B_iとする。

感想

  • Cとしては少し難しめかもしれないが、一歩ずつ思考を進めればちゃんと答えにたどり着ける良問だと思う。
  • 「差分の総和の最小化は中央値」ってのはわざわざ覚えるもんでもないと思う。
    • 上のように図を書けば分かる。
    • まぁその辺は個人のやり方があるので。