公式

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

Claude Opus 5 (High)

概要

各人について「自分より前にいる自分より背の高い人の数」\(A_x\) が与えられるので、元の並び \(P\) を復元する問題です。背の低い人から順に「空いている席のうち前から \(A_x+1\) 番目」に座らせていくことで復元でき、これを Fenwick 木(BIT)上の二分探索で高速に処理します。

考察

重要な観察:背の低い人から決めていく

\(A_x\) は「身長 \(x\) の人より前にいて、かつ身長が \(x\) より大きい人の数」です。ここで 身長 \(x\) より大きい人=身長 \(x+1, x+2, \dots, N\) の人 であることに注目します。

そこで、身長の小さい順 \(x = 1, 2, \dots, N\) に人を配置していくことを考えます。身長 \(x\) の人を置く瞬間には、

  • すでに席が決まっているのは身長 \(1, \dots, x-1\) の人(=\(x\) より低い人)だけ
  • まだ空いている席には、これから身長 \(x, x+1, \dots, N\) の人が座る

という状態です。つまり \(x\) より前にある空席には、必ず \(x\) より背の高い人が座る ことになります(\(x\) 自身の席以外はすべて \(x+1\) 以上の人)。逆に \(x\) より前の埋まっている席には \(x\) より低い人しかいません。

したがって:

身長 \(x\) の人の位置より前にある「空席の個数」= \(A_x\)

つまり、身長 \(x\) の人は、その時点で空いている席のうち前から \((A_x+1)\) 番目に座る ということが確定します。制約 \(A_x \le N - x\) は、この時点で空席が \(N-x+1\) 個あることに対応しており、必ず配置可能です。

具体例

\(N = 4\), \(A = (2, 0, 1, 0)\) の場合:

手順 空席 何番目の空席か 配置
\(x=1\) \(\{1,2,3,4\}\) \(A_1+1 = 3\) 位置 \(3\)
\(x=2\) \(\{1,2,4\}\) \(A_2+1 = 1\) 位置 \(1\)
\(x=3\) \(\{2,4\}\) \(A_3+1 = 2\) 位置 \(4\)
\(x=4\) \(\{2\}\) \(A_4+1 = 1\) 位置 \(2\)

結果 \(P = (2, 4, 1, 3)\)。実際に検算すると \(A=(2,0,1,0)\) となり正しいです。

素朴な実装ではなぜ間に合わないか

「空席を配列(またはリスト)で持ち、先頭から \(A_x+1\) 個目を線形に探す」あるいは「Python のリストから pop して詰める」実装は 1 回あたり \(O(N)\) かかり、全体で \(O(N^2)\) となります。\(N = 2\times 10^5\) では \(4\times 10^{10}\) 規模の操作となり TLE です。

そこで「\(k\) 番目の空席を求めて、それを埋める」という操作を高速に行えるデータ構造が必要になります。

アルゴリズム

Fenwick 木(BIT)を使い、位置 \(i\) が空席なら \(1\)、埋まっていれば \(0\) を管理します。すると

  • 「前から \(k\) 番目の空席を求める」=「累積和が初めて \(k\) 以上になる位置を求める」

という問題になり、これは BIT 上の二分探索(binary lifting) で \(O(\log N)\) で処理できます。

手順は次の通りです。

  1. BIT を「全要素 \(1\)」で初期化する。
  2. \(x = 1, 2, \dots, N\) について
    • \(k = A_x + 1\) として、累積和が \(k\) 以上になる最小の位置 \(p\) を BIT 上二分探索で求める。
    • \(P_p \leftarrow x\)。
    • 位置 \(p\) に \(-1\) を加算する(席を埋める)。
  3. \(P\) を出力する。

BIT 上二分探索は、bit を \(2^{\lfloor \log_2 N\rfloor}\) から \(1\) まで半分にしながら、 「\(pos + bit \le N\) かつ tree[pos+bit] < rem」なら \(pos \mathrel{+}= bit\)、\(rem \mathrel{-}= \) tree[pos+bit] と進める、という定型処理です。ループ終了時の \(pos\) は「累積和が \(rem\) 未満である最大の接頭辞長」なので、求める位置は \(1\)-indexed で \(pos+1\)、\(0\)-indexed 配列では添字 \(pos\) になります。

計算量

  • 時間計算量: \(O(N \log N)\)(各人につき BIT 上二分探索と 1 点更新が \(O(\log N)\))
  • 空間計算量: \(O(N)\)

実装のポイント

  • BIT の全 1 初期化を \(O(N)\) で行う:tree[i] = i & (-i)(\(i\) の最下位ビット)とするだけで、全要素 \(1\) の Fenwick 木が作れます。add を \(N\) 回呼ぶ \(O(N\log N)\) より簡潔かつ高速です。

  • \(0\)-indexed / \(1\)-indexed の混同に注意:BIT は \(1\)-indexed で扱い、出力配列 P は \(0\)-indexed にしています。上記の通り、二分探索の結果 \(pos\) をそのまま P[pos] に使うと辻褄が合います。

  • 入力の高速化:\(N = 2\times 10^5\) なので Python では sys.stdin.buffer.read().split() でまとめて読み、出力も ' '.join で 1 回にまとめると安全です。

  • 「\(A_x+1\) 番目の空席」の \(+1\) を忘れないこと(\(A_x = 0\) なら最も前の空席)。

  • 発想の別解として、逆に身長の大きい順に処理する方法や、平衡二分探索木・セグメント木上の下降で \(k\) 番目の空席を求める方法もありますが、本質は同じです。

    ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0])
    A = list(map(int, data[1:1 + n]))
    tree = [0] * (n + 1)
    for i in range(1, n + 1):
        tree[i] = i & (-i)
    LOG = n.bit_length()
    top = 1 << LOG
    P = [0] * n
    for x in range(1, n + 1):
        rem = A[x - 1] + 1
        pos = 0
        bit = top
        while bit:
            nx = pos + bit
            if nx <= n:
                t = tree[nx]
                if t < rem:
                    pos = nx
                    rem -= t
            bit >>= 1
        P[pos] = x
        i = pos + 1
        while i <= n:
            tree[i] -= 1
            i += i & (-i)
    sys.stdout.write(' '.join(map(str, P)))

main()

この解説は claude5opus-high によって生成されました。

投稿日時:
最終更新: