けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 005 B - Minimum Sum (400 点)

教育的な典型問題!

問題へのリンク

問題概要

 N 要素からなる順列  a が与えられる。 a の連続する各区間について、区間内の最小値を求め、その総和を求めよ。

制約

  •  1 \le N \le 2 × 10^{5}

考えたこと

まず思うのが、区間の個数は  O(N^{2}) 個ある。たとえ各区間に対する最小値を  O(1) で答えられたとしてもすべて求めて足すのでは間に合わない。

少し慣れればこの手の問題でかんがえるべきことは一瞬!!!

  • 数列の各要素  a_i について、それが合計される回数が何回ずつかを求める

こういう風に、全体の動きを追うのではなく、個別の要素に着目するのは典型な気がする。この問題はその典型成分をすごくシンプルな形で含んでいて教育的!!!

さて、要素  a_i区間 [  l, r) 内の最小値であるような  l, r の個数は、

  •  a_j <  a_i となるような最小の  j (>  i)
  •  a_j <  a_i となるような最小の  j (<  i)

がわかれば求められる。

その求め方

ここまで来て結構悩んだのだけど、結局僕は「ヒストグラム長方形面積最小値問題」で登場するようなスタックの使い方を思い出して実装した。

でも、

  •  a_i の値が小さい順に、その index i を set に追加して行って、 i より小さいところと大きいところを二分探索で求める

という風にすればよかった。こういう set の使い方に著しく不慣れなのが弱点。

とは言え、スタックで解いたこと自体は、計算量も落ちるし悪く無い。 O(N) で解けている。

#include <iostream>
#include <vector>
#include <stack>
using namespace std;
using pint = pair<int,int>;

int main() {
    int N; cin >> N;
    vector<long long> P(N);
    for (int i = 0; i < N; ++i) cin >> P[i];

    // prev[i] := i より手前で初めて P[i] より小さい値になる場所
    // next[i] := i より先で初めて P[i] より小さい値になる場所
    vector<int> prev(N, 0), next(N, N);
    stack<pint> st;
    st.push({0, -1}); // 番兵
    for (int i = 0; i < N; ++i) {
        while (st.size() > 0 && st.top().first >= P[i]) st.pop();
        prev[i] = st.top().second + 1;       
        st.push({P[i],i});
    }
    while (!st.empty()) st.pop();
    st.push({0, N});
    for (int i = N-1; i >= 0; --i) {
        while (st.size() > 0 && st.top().first >= P[i]) st.pop();
        next[i] = st.top().second;      
        st.push({P[i],i});
    }

    // 集計
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        res += P[i] * (next[i] - i) * (i+1 - prev[i]);
    }
    cout << res << endl;
}