Official

E - 列の復元 / Restoring the Sequence Editorial by admin

gpt-6-astra-medium

概要

各身長 \(x\) について「自分より前にいる、自分より背の高い人の人数」\(A_x\) から、元の列を復元します。

背の低い人から順に、前から \(A_x+1\) 番目の空き位置に配置することで復元でき、空き位置の管理には Fenwick Tree(BIT)を使います。

考察

背の低い人から配置する

最初に、答えの列に \(N\) 個の空き位置を用意します。そして、身長 \(1,2,\dots,N\) の順に人を配置することを考えます。

身長 \(x\) の人を配置する時点では、

  • 配置済みの人の身長は、すべて \(x\) より低い
  • \(x\) を配置したあと、残りの空き位置には、すべて \(x\) より高い人が入る

という状態です。

したがって、身長 \(x\) の人より前にいる背の高い人の人数は、\(x\) を配置する位置より前に残っている空き位置の個数と一致します。

この個数を \(A_x\) にするには、身長 \(x\) の人を 前から \(A_x+1\) 番目の空き位置に置けばよいことが分かります。

具体例

\(N=4\)、\(A=(2,0,1,0)\) の場合を考えます。空き位置を □ と書くと、次のように配置できます。

配置する身長 \(x\) \(A_x\) 選ぶ空き位置 配置後の列
\(1\) \(2\) 前から \(3\) 番目 □ □ 1 □
\(2\) \(0\) 前から \(1\) 番目 2 □ 1 □
\(3\) \(1\) 前から \(2\) 番目 2 □ 1 3
\(4\) \(0\) 前から \(1\) 番目 2 4 1 3

例えば、身長 \(1\) の人より前には身長 \(2,4\) の人がいるので、記録どおり \(A_1=2\) となっています。

空き位置を高速に探す

毎回、列を先頭から走査して空き位置を数えると、\(1\) 人の配置に最大 \(O(N)\) 時間かかります。全体では \(O(N^2)\) となり、\(N \le 2\times 10^5\) には間に合いません。

そこで、各位置について次の値を Fenwick Tree で管理します。

  • 空いている位置:\(1\)
  • 配置済みの位置:\(0\)

このとき、先頭からある位置までの和は、その範囲の空き位置の個数です。Fenwick Tree 上の二分累乗法を使うことで、前から \(k\) 番目の空き位置を \(O(\log N)\) 時間で探せます。

アルゴリズム

  1. 答えを格納する配列 ans を用意する。
  2. すべての位置が空いている状態、つまり各位置の値が \(1\) の Fenwick Tree を作る。
  3. \(x=1,2,\dots,N\) の順に、次の処理を行う。
    1. \(k=A_x+1\) とする。
    2. Fenwick Tree で、前から \(k\) 番目の空き位置を探す。
    3. その位置に身長 \(x\) の人を配置する。
    4. その位置の値を \(1\) から \(0\) に変更する。
  4. ans を出力する。

Fenwick Tree で \(k\) 番目の空き位置を探す方法

求める位置は、「先頭からの和が初めて \(k\) 以上になる位置」です。

コードでは、その直前にあたる「先頭からの和が \(k\) 未満となる最大の位置」を pos として探しています。

step を \(N\) 以下の最大の \(2\) のべき乗から始め、順に半分にしながら、次の処理を行います。

  • nxt = pos + step を候補とする。
  • nxt <= n かつ bit[nxt] < k なら、その区間をすべて飛ばしても目的の空き位置には届かない。
  • その場合は pos = nxt とし、飛ばした区間の空き数を k から引く。

この探索順では、bit[nxt] が飛ばそうとしている区間の空き数を表します。探索終了後、目的の位置は \(1\) 始まりで pos + 1 です。

正しさ

身長 \(x\) の人を配置するとき、前から \(A_x+1\) 番目の空き位置を選ぶため、その前にはちょうど \(A_x\) 個の空き位置があります。

それらには後で必ず \(x\) より高い人が配置されます。一方、すでに配置済みの人はすべて \(x\) より低いため、\(A_x\) の人数には含まれません。

よって、最終的に身長 \(x\) の人より前にいる背の高い人は、ちょうど \(A_x\) 人になります。これがすべての \(x\) について成り立つので、条件を満たす列が復元できます。

また、配置時の空き位置は \(N-x+1\) 個あり、制約から \(1 \le A_x+1 \le N-x+1\) なので、選ぶ位置は必ず存在します。

計算量

  • 時間計算量: \(O(N\log N)\)
  • 空間計算量: \(O(N)\)

各人について、空き位置の探索と Fenwick Tree の更新をそれぞれ \(O(\log N)\) 時間で行います。

実装のポイント

  • Fenwick Tree を \(O(N)\) で初期化する
    bit[i] は、末尾が位置 \(i\) で長さが i & -i の区間の和を持ちます。初期状態では各位置の値がすべて \(1\) なので、次のように直接初期化できます。

    bit = [i & -i for i in range(n + 1)]
    
  • 添字の違いに注意する
    Fenwick Tree は \(1\) 始まり、ans は \(0\) 始まりです。探索後の目的の位置は Fenwick Tree 上では pos + 1 なので、答えへの代入は ans[pos] = x となります。

  • 探索条件は bit[nxt] < k
    bit[nxt] == k の場合、その区間内に目的の空き位置があります。区間全体を飛ばしてはいけません。

  • 配置した位置は \(1\) 減らす
    使用した位置を空き位置として数えないよう、Fenwick Tree の該当位置に \(-1\) を加算します。

    ソースコード

import sys


def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    n = data[0]

    bit = [i & -i for i in range(n + 1)]
    ans = [0] * n
    highest_bit = 1 << (n.bit_length() - 1)

    for x in range(1, n + 1):
        k = data[x] + 1
        pos = 0
        step = highest_bit

        while step:
            nxt = pos + step
            if nxt <= n and bit[nxt] < k:
                k -= bit[nxt]
                pos = nxt
            step >>= 1

        ans[pos] = x
        idx = pos + 1
        while idx <= n:
            bit[idx] -= 1
            idx += idx & -idx

    sys.stdout.write(" ".join(map(str, ans)) + "\n")


if __name__ == "__main__":
    main()

この解説は gpt-6-astra-medium によって生成されました。

posted:
last update: