2025-01-04

AtCoder ARC 189 D - Takahashi is Slime (3D, 黄色, 700 点)

$O(N)$ で解けた！　同じ値が連続する場合の対処に苦慮した。

問題へのリンク

問題概要

強さが $A_{1}, A_{2}, \dots, A_{N}$ のスライムが一列にこの順に並んでいる。 $k = 1, 2, \dots, N$ に対して、次の問に答えよ。

【問】
初期状態で左から $K$ 番目にいるスライムが高橋君である。高橋君が下記の行動を好きな回数だけ行ったあとの、高橋君の大きさとしてあり得る最大値を求めたい。

高橋君に隣接するスライムのうち、高橋君より大きさが真に小さいものを選んで吸収する
その結果、吸収されたスライムは消滅し、高橋君の大きさは吸収したスライムの大きさだけ増加する
上記の操作の際、スライムが吸収され消滅したことで生じた隙間は直ちに埋められ、消滅したスライムの両端のスライム同士は新たに隣接する

制約

$2 \le N \le 5 \times 10^{5}$
$1 \le A_{i} \le 10^{9}$

考えたこと

まず最初に考えたのは、「強さが最大のスライムは全スライムを吸収できるだろう」ということだった（のちに、最強のスライムが複数体いて連続して並んでいると反例があることがわかった）。

そして、強さが最大のスライムの index を $i$ としたとき、その左側の区間 $\lbrack 0, i)$ に関する問題と、右側の区間 $\lbrack i+1, N)$ に関する問題に分割して解けるだろうと考えた。それ以降は、再帰的に、区間内の最強のスライムで左右を分割して解けるだろう。

区間を再帰的に分割して解く部分については、左右それぞれの区間について、次のようにすればよいだろうと考えた。ただし、簡単のため、一旦、最強のスライムは単独で存在するとする。

区間 $\lbrack l, r)$ について解くことを考える。区間 $\lbrack l, r)$ が切り出される直前の最強のスライムを $p$ とする。

区間 $\lbrack l, r)$ で最強のスライムを考える（index を $m$ とする）
区間 $\lbrack l, r)$ に含まれるスライムの強さの総和を $S$ として
のとき：
- 直前の最強のスライム $p$ さえ吸収できる
- よって、スライム $m$ の答えは、スライム $p$ の答えに一致する
のとき：
- 直前の最強のスライム $p$ は吸収できない
- よって、スライム $m$ の答えは、区間 $\lbrack l, r)$ から広げることはできないから、 $S$ である
これ以降、再帰的に左側の区間 $\lbrack l, m)$ の問題と、右側の区間 $\lbrack m+1, r)$ の問題とに分割して考える

これで意気揚々と解けたと思ったが、区間内で最強のスライムが複数体いるときが面倒だと気づいた。

区間内で最強のスライムが複数体いるとき

まず、最強のスライムが連続して並んでいない場合は気にしなくてよい。連続して並ぶ場合が扱いづらい。たとえば、

8
3 15 15 15 15 15 7 8

の場合の答えは

3 93 15 15 15 93 7 15

となる。つまり、最強のスライムが 5 個並んでいるが、次のように整理できる。

その両端のスライムに関しては、区間内のスライムをすべて吸収できる
その内側のスライムに関しては、どのスライムも吸収できない

先ほどの分割統治法の中に組み込む際には、この「連続する最強のスライム」については、ひとまとまりにして扱うこととした。つまり、連続する最強のスライムはすべてまとめて処理してしまってから、その左右の区間に分割して解いていくことにした。

計算量

この解法では、Cartesian 木がぴったりだ。次の記事に詳しく書いた。

drken1215.hatenablog.com

Cartesian 木を使うと「区間内の最強のスライムで左右に分岐していく」処理を自然に書ける。そして、全体として $O(N)$ の計算量で実現できる。

コード

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

// Cartesian Tree
template<class T> struct CartesianTree {
    int root;  // root
    vector<int> par, left, right;

    CartesianTree() {}
    CartesianTree(const vector<T>& v) : root(0)
    , par(v.size(), -1), left(v.size(), -1), right(v.size(), -1) {
        vector<int> st(v.size(), 0);
        int top = 0;
        for (int i = 1; i < v.size(); ++i) {
            if (v[st[top]] > v[i]) {
                while (top >= 1 && v[st[top - 1]] > v[i]) --top;
                par[left[i] = st[top]] = i;
                if (top == 0) root = i;
                else right[par[i] = st[top - 1]] = i;
                st[top] = i;
            } else {
                right[par[i] = st[top]] = i;
                st[++top] = i;
            }
        }
    }
};

int main() {
    int N;
    cin >> N;
    vector<long long> A(N), B(N), S(N+1, 0), res(N, -1);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
        B[i] = -A[i];  // Cartesian 木を大きいものから取るように変更
        S[i+1] = S[i] + A[i];
    }

    // Cartesian 木を用いて再帰的に探索
    CartesianTree<long long> ct(B);
    auto rec = [&](auto rec, int left, int right, int v) -> void {
        if (right - left < 1) return;

        // 同じ値が続く範囲を求める
        int vr = v;
        while (vr < right && A[vr] == A[v]) vr++;
        
        for (int x = v; x < vr; x++) {
            // A[x] の両隣に A[x] より小さいものがあるか
            bool exist = false;
            if (x-1 >= 0 && A[x-1] < A[x]) exist = true;
            if (x+1 < N && A[x+1] < A[x]) exist = true;

            if (exist) {
                long long sum = S[right] - S[left];
                if (ct.par[v] == -1 || sum <= A[ct.par[v]]) res[x] = sum;
                else res[x] = max(sum, res[ct.par[v]]);
            } else {
                res[x] = A[x];
            }
        }
        if (ct.left[v] != -1) rec(rec, left, v, ct.left[v]);
        if (ct.right[vr-1] != -1) rec(rec, vr, right, ct.right[vr-1]);
    };
    rec(rec, 0, N, ct.root);

    // 出力
    for (auto val : res) cout << val << " ";
    cout << endl;
}