【AtCoder】ABC 392 F - Insert | 緑コーダーが解くAtCoder

atcoder.jp

配点: 500 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 1395

問題概要
制約
考察
実装例
コメント

問題概要

空の配列 $A$ があり、 $i=1,2,\cdots,N$ の順に以下の操作を行う。

数 $i$ を、 $A$ の前から $P_i$ 番目の位置になるように挿入する。

全ての操作を終えた後の $A$ を出力せよ。

制約

$1 \leq N \leq 5\times 10^5$
$1 \leq P_i \leq i$
入力は全て整数。

考察

要素を配列の任意の位置に挿入するという操作は高速に処理することが難しく、本問の制約下では TLE してしまう。

とはいえ、ランダムアクセスをしようにも、与えられる数列 $P$ は挿入するインデックスを表すので、単純に前から決めていくことはできない。

しかし、一か所だけ位置がすぐに確定する場所がある。

数 $N$ の位置だ。これは、 $A_{P_N}$ である。

この要領で、操作を逆から考えてみよう。

$i = 1, 2, \cdots, N$ に対して、数 $i$ を $A$ の前から $P_i$ 番目の位置になるように挿入する。

$\Leftrightarrow$ $i = N, N-1, \cdots, 1$ に対して、数 $i$ を $A$ の未確定の要素の中で前から $P_i$ 番目に書き込む。

この操作の言い換えによって、目標は

$A$ の前から $x$ 番目の未確定要素を高速に取得する

ことになった。

これは、「未確定のインデックスを $0$ 、確定済みのインデックスを $1$ とした区間和 Segment Tree を用いて、 $\mathrm{prod}(0, r) \lt P_i$ を満たす最大の $r$ を二分探索する」ことによって実現可能である。

※ $\mathrm{prod}(0, r) :=$ 半開区間 $[0, r)$ の区間和

実装においては、 ACL の Segtree を使うのが楽だろう。先ほどの二分探索はmax_rightで可能だが、実装例では比較関数の定義をラムダ式を使って書いていることに注意。

実装例

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

#if __has_include(<atcoder/all>)
#include <atcoder/all>
using namespace atcoder;
#endif

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
#define rrep(i, start, end) for (auto i = (start); (i) >= (end); (i)--)

// ======================================== //

int op(int a, int b) { return a + b; }
int e() { return 0; }

int main()
{
    int N;
    cin >> N;
    vector<int> P(N);
    rep(i, 0, N) cin >> P[i];

    segtree<int, op, e> seg(vector<int>(N, 1));
    vector<int> ans(N);
    rrep(i, N - 1, 0)
    {
        auto searchPos = [&](int x)
        {
            auto f = [&](int v)
            { return v < x; };
            
            return seg.max_right(0, f);
        };

        int j = searchPos(P[i]);
        ans[j] = i + 1;

        seg.set(j, 0);
    }

    rep(i, 0, N) cout << ans[i] << ' ';
    cout << endl;

    return 0;
}