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\) について、この順に以下の操作を行なう。
- \(x \coloneqq \min(S)\) とする。 \(P_i > x\) ならば、\(S\) から \(x\) を削除し、代わりに \(P_i\) を挿入する。
- \(\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\) について、この順に以下の操作を行なう。
- \(P_i > x\) ならば、\(S\) から \(x\) を削除し、代わりに \(P_i\) を挿入する。
- \(x \notin S\) である間、 \(x \leftarrow x + 1\) とする。
- \(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: