Official

D - Prefix K-th Max Editorial by penguinman


問題文中で問われているのは \(P\) の先頭 \(i\) 項のうち \(K\) 番目に大きい値のみですが、ここでは \(i=K,K+1,\ldots,N\) について \(P\) の先頭 \(i\) 項のうち値の大きい方から \(K\) 個の要素からなる集合 \(S_i\) を求めに行くことを考えます。

このような \(N-K+1\) 個の集合を愚直に求めた場合、時間計算量が \(O(NK)\)\(O(NK \log K)\) となり実行時間制限に間に合いません。また \(N-K+1\) 個の集合をそれぞれ別個の配列等に保存した場合、空間計算量が \(O(NK)\) となってしまい実行時間制限、およびメモリ制限を超過してしまいます。

そのため、以下の発想が必要となります。

  • \(i\ (K \lt i)\) について、\(S_{i-1}\) の内容を利用して \(S_i\) を求める。保持するデータ構造も \(S_{i-1}\) のものを再利用する。

では、具体的にどのようにすればよいのでしょうか?

結論から言うと、以下のようなアルゴリズムを用いればよいです。

まず、\(S_K\)\(\{P_1,P_2,\ldots,P_K\}\) で初期化する。

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

  1. \(S_i\)\(S_{i-1}\) で初期化する(実際には保持するデータ構造を再利用するためこの操作は行わない)。
  2. \(S_{i-1}\) に含まれる要素のうち値が最小のものを \(x\) とする。\(x \lt P_i\) ならば、\(S_i\) から \(x\) を削除し、代わりに \(P_i\) を挿入する。

上記のアルゴリズムも、愚直に計算した場合 \(O(NK)\) の時間計算量となります。しかし、大抵のプログラミング言語の標準ライブラリに備わっている「優先度付きキュー」と呼ばれるデータ構造を用いると以下の操作をそれぞれ \(O(\log n)\ (n\) は操作時点での要素数\()\) で行うことができ、これにより時間計算量を \(O(N \log K)\) まで改善することが可能です。

  • 最小値の取り出し
  • 最小値の削除
  • 要素の挿入

求めるべき答えは、\(i=K,K+1,\ldots,N\) についての \(S_i\) に含まれる要素の最小値です。そのような要素の取り出しももちろん、\(1\) 回あたり対数時間で行うことが可能です。

実装例 (Python): Python においては「heapq」という名前で「優先度付きキュー」が実装されています。

import heapq
N,K = map(int,input().split())
P = list(map(int,input().split()))
que = P[0:K]
print(min(que))
heapq.heapify(que)
for i in range(K,N):
    minima = heapq.heappop(que)
    minima = max(minima,P[i])
    heapq.heappush(que,minima)
    ans = heapq.heappop(que)
    print(ans)
    heapq.heappush(que,ans)

実装例 (C++): C++ においては「priority_queue」という名前で「優先度付きキュー」が実装されています。C++ の場合、デフォルトでは最小値ではなく最大値が取り出されてしまうことに注意してください。

#include<bits/stdc++.h>
using namespace std;

int main(){
	int N,K; cin >> N >> K;
	vector<int> P(N);
	for(int i=0; i<N; i++) cin >> P[i];
	priority_queue<int,vector<int>,greater<int> > que;
	for(int i=0; i<K; i++) que.push(P[i]);
	cout << que.top() << endl;
	for(int i=K; i<N; i++){
		if(que.top() < P[i]){
			que.pop();
			que.push(P[i]);
		}
		cout << que.top() << endl;
	}
}

posted:
last update: