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

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

AtCoder ABC 359 E - Water Tank (1D, 水色, 500 点)

AtCoder AtCoder500点 ABC-E 水色diff NoviSteps1D DP高速化:stackの活用 stack シミュレーション:stackの活用 DP シミュレーション 所要時間を求める問題 各kに対して

いかにも stack が登場しそうな問題！

問題概要

下の図のように、高さが H_{1}, H_{2}, \dots, H_{N} の仕切りが等間隔に並んでいて、その間と両端の N+1 カ所のスペースを表す番号を 0, 1, \dots, N とする。これらは幅が等しい。

今、スペース 0 に水を入れていく。1 個分のスペースについて、水位は 1 秒で 1 だけ上昇する量の水である。

このとき、各 i = 1, 2, \dots, N について、スペース i に水が入り始める瞬間の時刻 + 1 の値を求めよ。

制約

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

考えたこと

スペース i に水が入り始めるのは、スペース i-1 の水位が H_{i} になった瞬間である。その瞬間の水の量を求めればよい。今、次のように配列 dp を定義しよう。

dp[i] ← スペース i に水が入り始める瞬間の時刻

すると、下図のように H_{j} \lt H_{i} となる最大の [tex j] ( j \lt i) をとったとき、

dp[i] = dp[j] + H[i] * (i - j)

と表せる。なお、ここで H_{0} = \infty としておく。

あとは、各 i に対して、 H_{j} \lt H_{i} となる最大の [tex j] ( j \lt i) を求める必要がある。これはよく知られた有名問題で、stack を使って O(N) の計算量で求められる。詳しくは次の問題の解説（鉄則本にある）を参照。

drken1215.hatenablog.com

全体として配列 dp を求める計算量は O(N) と評価できる。

コード

#include <bits/stdc++.h>
using namespace std;
const long long INF = 1LL<<60;

int main() {
    long long N;
    cin >> N;
    vector<long long> H(N + 1, INF);
    for (int i = 1; i <= N; i++) cin >> H[i];

    // H[j] > H[i] となる最大の j を求める
    vector<long long> prev(N+ 1);
    vector<pair<long long, int>> stac;
    for (int i = 0; i <= N; i++) {
        while (!stac.empty() && stac.back().first <= H[i]) stac.pop_back();
        if (stac.empty()) prev[i] = 0;
        else prev[i] = stac.back().second;
        stac.emplace_back(H[i], i);
    }

    // DP
    vector<long long> dp(N+1, 0);
    for (int i = 1; i <= N; i++) {
        dp[i] = dp[prev[i]] + H[i] * (i - prev[i]);
        cout << dp[i]+1 << " ";
    }
    cout << endl;
}