AtCoder Educational DP Contest Q問題解いてみた

今回の問題

問題文

$N$ 本の花が横一列に並んでいます。
各 $i$ ($1 \leq i \leq N$) について、左から $i$ 番目の花の高さは $h_i$ で、美しさは $a_i$ です。
ただし、$h_1, h_2, \ldots, h_N$ はすべて相異なります。

太郎君は何本かの花を抜き去ることで、次の条件が成り立つようにしようとしています。

• 残りの花を左から順に見ると、高さが単調増加になっている。

残りの花の美しさの総和の最大値を求めてください。

制約

• 入力はすべて整数である。
• $1 \leq N \leq 2 \times 10^5$
• $1 \leq h_i \leq N$
• $h_i \neq h_j$ ($i \neq j$)
• $1 \leq a_i \leq 10^9$

回答

#include <bits/stdc++.h>
using namespace std;

// セグメント木の実装 (Range Maximum Query)
class SegmentTree {
public:
    int n;
    vector<long long> tree;

    // コンストラクタ
    SegmentTree(int size) {
        n = 1;
        while (n < size) n *= 2;
        tree.assign(2 * n, 0);
    }

    // 点更新: index の値を val に更新 (今回は最大値なので max をとる)
    void update(int index, long long val) {
        index += n;
        tree[index] = val;
        while (index > 1) {
            index /= 2;
            tree[index] = max(tree[2 * index], tree[2 * index + 1]);
        }
    }

    // 区間取得: [l, r) の最大値を返す
    long long query(int l, int r) {
        l += n;
        r += n;
        long long res = 0;
        while (l < r) {
            if (l & 1) res = max(res, tree[l++]);
            if (r & 1) res = max(res, tree[--r]);
            l /= 2;
            r /= 2;
        }
        return res;
    }
};

int main() {
    int N;
    cin >> N;

    vector<int> h(N);
    for (int i = 0; i < N; i++) cin >> h[i];

    vector<int> a(N);
    for (int i = 0; i < N; i++) cin >> a[i];

    // 高さが最大 N まであるので、サイズ N+1 のセグメント木を作成
    SegmentTree seg(N + 1);

    long long ans = 0;

    for (int i = 0; i < N; i++) {
        // 条件: 自分より左にあって、高さが自分(h[i])より低い花の DP の最大値を取得
        // 区間 [0, h[i]) の最大値
        long long max_prev = seg.query(0, h[i]);
        
        // 現在の花を選んだ場合の総和
        long long current_val = max_prev + a[i];

        // セグメント木の 位置 h[i] を更新
        seg.update(h[i], current_val);
        
        // 全体を通しての最大値を更新
        ans = max(ans, current_val);
    }

    cout << ans << endl;

    return 0;
}

自分なりの解説

dp[i]をi番目を選んだ時の美しさの最大値とするというところまでは自力で分かったが、その遷移のためにはdp[i-1]番目までの高さがhi以下の最大値を見つけてこなくてはいけない。これを愚直にやると、実行時間オーバーしてしまう。調べたところセグメントツリーを使うことにより、区間の最大値を素早く更新、取得ができるのでそれを使う。