【AGC005】B - Minimum Sum

atcoder.jp

ABC214のD問題に似ていたので解きました。

解法

lとrを全探索していたらO(N²)になるため、間に合いません。
なので、個々の要素が答えにどれだけ寄与するか、という主客転倒の考え方を使ってときます。

小さい要素ほど、多くの区間に置いて最小値になりやすいです。
逆に大きい要素は最小値になりにくいです。
今回は、大きい要素から見ていって答えへの寄与を考えます。

今見ている要素のindexをjとします。
jの左の要素を見てその値がa[j]より大きい場合、辺で繋ぎます。
右の要素も同様にします。

サンプル3でj=3の場合、下のようなグラフが出来上がっているはずです。
大きい要素から見ていっているため、j=3は一番最後になります。

j=3が最小値になる区間は、l=0~3, r=3~7 の場合です。
これを満たす区間の数は、4 * 5 = 20になります。
ということでj=3の時、答えに与える影響は、 a[j] * 4 * 5であることがわかりました。
これを一般化すると、
j = kの要素が答えに与える影響は、 a[j] * (a[j-1]と連結の要素数 + 1) * (a[j+1]と連結の要素数 + 1 )
となります。

実装

連結する要素数を素早く求めたいので、Union-Find木を使うことにしました。

void solve(long long N, std::vector<long long> a){
    vector<pair<ll, ll>> A;
    REP (i, N) A.push_back({a[i], i});
    sort(rall(A));

    UnionFind Uf(N);
    ll ans = 0;
    REP (i, N) {
        int j = A[i].second;
        int l = j-1, r = j+1;
        int lsiz = 1, rsiz = 1;

        if (l >= 0 && a[l] >= a[j]) {
            lsiz += Uf.size(l);
            Uf.unite(l, j);
        }
        if (r < N && a[r] >= a[j]) {
            rsiz += Uf.size(r);
            Uf.unite(r, j);
        }

        ans += a[j] * lsiz * rsiz;
    }

    c(ans)
}