Ex - Min + Sum Editorial by MMNMM

ちょっとした別解

公式解説の方針を前提とします。

公式解説では、列を分割して再帰的に問題を解いています。 この方針では大きい問題を解くためにより小さい問題の結果を使っていないため、どの区間から解いても結果が変わりません。

公式解説の \(\lbrack2\rbrack\) では、Cartesian Tree の行きがけ順で問題を解いていっていると考えられます。 以下の実装例では Cartesian Tree の通りがけ順で問題を解いていると考えられます。

時間計算量は公式解説と同じ、\(O(N\log^2N)\) です。

#include <iostream>
#include <vector>
#include <numeric>
#include <utility>
#include <algorithm>

int main() {
    using namespace std;

    unsigned long N, S;
    cin >> N >> S;
    vector<unsigned long> A(N), B(N);
    for (auto &&a : A)cin >> a;
    for (auto &&b : B)cin >> b;

    // B の (suffix) 累積和
    // sum_B[i] = B[i] + B[i + 1] + ... + B[N - 1];
    vector<unsigned long> sum_B{B};
    sum_B.emplace_back();
    inclusive_scan(rbegin(sum_B), rend(sum_B), rbegin(sum_B));

    // A[i] に対する (L, R) を求める
    // O(N) 時間
    vector<pair<unsigned long, unsigned long>> A_interval(N, {0, N});
    {
        vector<unsigned long> idx;
        for (unsigned long i{}; i < N; ++i) {
            while (!empty(idx) && A[idx.back()] > A[i]) {
                A_interval[idx.back()].second = i;
                idx.pop_back();
            }
            if (!empty(idx)) A_interval[i].first = idx.back() + 1;
            idx.emplace_back(i);
        }
    }

    unsigned long ans{};
    for (unsigned long i{}; i < N; ++i) {
        // i ごとに問題を解く
        const auto&[L, R]{A_interval[i]};
        const auto first{next(begin(sum_B), L)};
        const auto mid{next(begin(sum_B), i + 1)};
        const auto last{next(begin(sum_B), R + 1)};
        // i - L(i) + 1 と R(i) - i の大きくないほうを固定
        if (2 * i < L + R)
            for (unsigned long l{L}; l <= i; ++l)
                // A[i] + sum_B[l] - S 以上の sum_B[r] の個数
                ans += upper_bound(mid, last, max(A[i] + sum_B[l], S) - S, greater<>{}) - mid;
        else
            for (unsigned long r{i + 1}; r <= R; ++r)
                // S + sum_B[r] - A[i] 以下の sum_B[l] の個数
                ans += mid - lower_bound(first, mid, max(A[i], S + sum_B[r]) - A[i], greater<>{});
    }
    cout << ans << endl;

    return 0;
}

posted:
last update: