Mister雑記

競プロやります。

ARC100-D 「Equal Cut」 (600)

分かってしまえばそれしかないが、なんで思いついたんだろう。

問題リンク

概要

長さ Nの数列 Aを3箇所で切って出来た4つの空でない数列について、それぞれの総和を P, Q, R, Sとする。このとき、「 P, Q, R, Sの最大値と最小値の差」の最小値を求めよ。

制約

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

解法

まず何となく「 P, Q, R, S Aの総和を大体4等分するんだろうな〜」と考えて、

  • 和が Aの総和の1/4になるまで数列を伸ばして、超えたらその前後で切る

みたいな解法を思いつく。しかしそれではサンプルにあるケースがどうにも上手くいかないことに気づき、ここで行き詰まる。

次に注目したのが問題文。いわゆるメタ解きというやつ。

  • なぜカット数は変数ではなく3で固定なのか?
  •  Nの制約的に解法は O(N \log N)くらいだろう

の2点に注目した結果、

  1. まず Aの2つの切り方を全探索する。
  2. 各パートについて、なるべく2等分するような切断位置を二分探索で探す。

という解法に至った。これなら O(N \log N)で間に合うし、先程考えた解法は4つだと上手くいかないが2つなら上手くいくことも証明できるのでここで実装に移る。

二分探索する上でカットしたときの数列の和を O(1)で求める必要があるが、これは累積和を作っておけば大丈夫。

あと注意点として、数列をなるべく均等に2つに切ったときのそれぞれの和を前から P, Qとしたとき、 P \leq Q P \gt Qの2パターンを考える必要がある。これは二分探索で境界となった値をどっちに入れるかを実際に試せばいい。

実装例

#include <algorithm>
#include <iostream>

using namespace std;

using ll = long long;

ll N;
ll A[200010], sum[200010];
// 入力の数列と、その累積和(1-indexed)

// (A_b, ..., A_e)を大体半分にカットし、前半と後半の総和を返す
pair<ll, ll> divide(ll b, ll e) {
    ll ok = b, ng = e;
    ll ds = sum[e] - sum[b - 1];

    // 二分探索で真ん中ちょい手前を探る
    while (ng - ok > 1) {
        ll mid = (ok + ng) / 2;
        ll ps = sum[mid] - sum[b - 1];
        if (ps * 2 <= ds) {
            ok = mid;
        } else {
            ng = mid;
        }
    }

    // ちょっと手前で切った場合の値
    pair<ll, ll> former = make_pair(sum[ok] - sum[b - 1], sum[e] - sum[ok]);

    // ちょっと後ろで切った場合の値
    pair<ll, ll> later = make_pair(sum[ok + 1] - sum[b - 1], sum[e] - sum[ok + 1]);

    // より差の小さい方を選ぶ
    if (abs(former.first - former.second) < abs(later.first - later.second)) {
        return former;
    } else {
        return later;
    }
}

// Aを[0, itr]と(itr, N)で切ったときの解を返す
ll cut(ll itr) {
    pair<ll, ll> res[2] = {divide(1, itr), divide(itr + 1, N)};

    ll elem[4];
    for (ll i = 0; i < 2; ++i) {
        elem[i * 2] = res[i].first;
        elem[i * 2 + 1] = res[i].second;
    }

    sort(elem, elem + 4);
    return elem[3] - elem[0];
}

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

    ll ans = sum[N];
    for (ll i = 2; i < N; ++i) {
        ans = min(ans, cut(i));
    }

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

前の記事のときよりは無駄を省いた。それでもやはり実装はやや重め。

ちなみにeditorialは二分探索ではなく尺取り法を使っている。にぶたんと尺取りは生き別れの兄弟なので(ん?)、総和系の場合はどちらでも実装できることが多い。

感想

  • 冒頭でも言ったが、この解法に気づいたのは単なる幸運に近い。
  • 「もしかして数列系の問題が得意なのでは?」と思うようになったキッカケの問題でもある。