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)\) 時間で探せます。
アルゴリズム
- 答えを格納する配列
ansを用意する。 - すべての位置が空いている状態、つまり各位置の値が \(1\) の Fenwick Tree を作る。
- \(x=1,2,\dots,N\) の順に、次の処理を行う。
- \(k=A_x+1\) とする。
- Fenwick Tree で、前から \(k\) 番目の空き位置を探す。
- その位置に身長 \(x\) の人を配置する。
- その位置の値を \(1\) から \(0\) に変更する。
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: