D - Prefix K-th Max Editorial by Forested


公式解説は \(\Theta(N \log K)\) 時間での解法ですが、少し工夫をすると \(\Theta(N)\) 時間にすることが可能です。

公式解説のアルゴリズムは以下のようなものでした。

まず、 \(S \coloneqq \lbrace P_1, P_2, \ldots, P_K \rbrace\) とする。 \(\min(S)\) を出力する。

その後、 \(i = K+1, K+2, \ldots, N\) について、この順に以下の操作を行なう。

  1. \(x \coloneqq \min(S)\) とする。 \(P_i > x\) ならば、\(S\) から \(x\) を削除し、代わりに \(P_i\) を挿入する。
  2. \(\min(S)\) を出力する。

このアルゴリズムでは、 \(x\) の値は前ステップの値よりも小さくなることはありません。よって、前ステップでの \(x\) の値を記憶しておくことで高速化ができます。具体的には、次のようなアルゴリズムで \(\Theta(N)\) 時間が達成できます。

まず、 \(S \coloneqq \lbrace P_1, P_2, \ldots, P_K \rbrace, x \coloneqq \min(S)\) とする。 \(x\) を出力する。

その後、 \(i = K+1, K+2, \ldots, N\) について、この順に以下の操作を行なう。

  1. \(P_i > x\) ならば、\(S\) から \(x\) を削除し、代わりに \(P_i\) を挿入する。
  2. \(x \notin S\) である間、 \(x \leftarrow x + 1\) とする。
  3. \(x\) を出力する。

\(S\) の管理ですが、各要素は \(1\) 以上 \(N\) 以下ですから、 \(A[k] \coloneqq (k \in S)\) として長さ \(N\) の真偽値の配列を管理すれば良いです。

実装例 (Python)

n, k = map(int, input().split())
p = list(map(lambda x: int(x) - 1, input().split()))

a = [False] * n
for i in range(k):
    a[p[i]] = True
x = min(p[:k])
print(x + 1)
for i in range(k, n):
    if p[i] > x:
        a[x] = False
        a[p[i]] = True
    while not a[x]:
        x += 1
    print(x + 1)

posted:
last update: