AtCoder: ABC140-E Second Sum (500)

問題はこちら

atcoder.jp

長さ の順列 が与えられた時、区間 に対して以下を計算せよ。

• 全ての重複しない区間において、二番目に大きい数の和。

考え方

簡単バージョンとしてこの問題があるので見ると良い。

atcoder.jp

step1 - 問題の言い換え

愚直にやると、全ての区間を列挙して、その区間の二番目の要素を足していくことになるが、これは[tex: N³] のオーダーとなり到底間に合わない。そこで、まずはそれぞれの要素が何回足されることになるかを考えてみる。（区間系の問題では、それぞれの要素が何回操作されるかを考えてみると良いことが多い気がする。）

が足される回数は、以下のような場合である。

• の右側に よりも大きいものが一つあり、左側には よりも小さいもののみある。
• 上の逆。

よって、上の図のように を定めると、一つの が答えに寄与する回数は

step2 - 実装方法

問題は、これをどのようにしてTLEしないように求めるかである。愚直にそれぞれの要素を見て、自分よりも大きい要素を探しに行くと かかるので間に合わない。

どうにかしてlower_bound()などを使って、高速に自分より大きい要素を探したいが、このままだとindexの情報を保ったまま探索をすることができない。

ここでは の要素を大きい順に見ていき、すでに見た要素のindexをsetに入れて管理する。こうすることで、今見ている要素よりも大きい要素のみのindexが入っており、lower_bound() などで高速に今見ている要素よりも大きく、かつ一番右に近いものを見つけることができる。これが見つかれば、あとは境界条件に気をつけてイテレーターを動かして二番目の要素までを得る。

また、今回の問題では、個数を数えるときに端を含まないように数えなければならない（つまり、二番目に大きい要素を含めないようにする）ので、 は で初期化し、 は で初期化することで、うまく個数を数えることができる。

解答

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using vl = vector<ll>;

#define rep(i, n) for(ll i = 0; i < (ll)(n); i++)

int main() {

    ll n;
    cin >> n;
    vl p(n);
    unordered_map<ll, ll> mp;
    set<ll> s;
    rep(i, n) {
        cin >> p[i];
        mp[p[i]] = i;
        s.insert(p[i]);
    }

    set<ll> idx;
    ll ans = 0;
    while (!s.empty()) {
        ll now = *s.rbegin();
        s.erase(now);

        ll i = mp[now];
        vl l(2, -1), r(2, n);
        {
            auto itr = idx.lower_bound(i);
            rep(j, 2) { // r
                if (itr == idx.end()) break;
                r[j] = *itr;
                itr++;
            }
        }
        {
            auto itr = idx.lower_bound(i);
            rep(j, 2) { // l
                if (itr == idx.begin()) break;
                itr--;
                l[j] = *itr;
            }
        }

        ll l1 = i - l[0], l2 = l[0] - l[1];
        ll r1 = r[0] - i, r2 = r[1] - r[0];
        ans += (l1 * r2 + l2 * r1) * now;

        idx.insert(i);
    }
    cout << ans << '\n';

    return 0;
}

Submission #7453131 - AtCoder Beginner Contest 140