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)\) で処理できます。
手順は次の通りです。
- BIT を「全要素 \(1\)」で初期化する。
- \(x = 1, 2, \dots, N\) について
- \(k = A_x + 1\) として、累積和が \(k\) 以上になる最小の位置 \(p\) を BIT 上二分探索で求める。
- \(P_p \leftarrow x\)。
- 位置 \(p\) に \(-1\) を加算する(席を埋める)。
- \(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 によって生成されました。
投稿日時:
最終更新: