【ABC234】D - Prefix K-th MaxをC++で解説（解法２つあり）

はじめに

　皆さんこんにちは！ 競プロの問題の理解を深めるために解説記事を書いていこうと思います！今回は ABC234 D です！

気に入っていただけたらLGTMして頂けるとうれしいです！

わからないことなどあれば下のコメントやTwitterからお願いいたします!

D - Prefix K-th Max

問題の詳細は上の見出しからいけます．
Difficulty : 503

問題概要

$(1, 2, ..., N)$の順列$P = (P_1, P_2, ..., P_N)$において正整数$K$が与えられたとき, $i = K, K+1, ..., N$について
$P$の先頭$i$のうち$K$番目に大きい値を求めよ．

制約

・$1 \leq K \leq N \leq 5 × 10^5$
・$(P_1, P_2, ..., P_N)$は$(1, 2, ..., N)$の並び替えによって得られる
・入力はすべて正整数

考察

ポイント

K番目に大きい数を出力できる構造を使おう！

解説

入力が$K$番目から出力が必要になります．
入力毎にソートして$K$番目に大きい値を求めようとすると計算量は$O(NKlogK)$となりTLEになります．計算量を減らすことを考えましょう．
$K$番目に大きい値を求めるので毎回データを残すのは$K$個のデータで十分です．
　入力$P_i (i = K+1, ... , N)$での$K$番目の値を$a_i$とします．
　$a_{i+1}$は$a_{i+1} < a_i$のとき**$a_{i+1}$は$K$番目より後の大きさの値となるので出力は$a_i$となります．
　逆に$a_{i+1} > a_i$のとき$a_{i+1}$は$1 ～ K$番目になるので，出力は$a_i$より大きい値のうち次に大きい値となります．
　$K$個のデータを順番通りに管理するときに優先度付きキュー**を用いるか，setを用いて二分探索することにより計算量を$O(NlogK)$にすることができます．

　それでは実装していきましょう．

実装

実装1 ~優先度付きキュー(priority_queue)~

$K-1$番目まではpriority_queueにpushします．

一番最初の出力はtop(1番前の値を取り出す)した値を出力し，それ以降は入力がtopした値より大きい場合，**topした値をpop(priority_queueから消去)**します．
　そして入力をpriority_queueにpopして，出力はtopの値になります．
　
　それ以降は入力がtopした値より小さい場合，出力はtopの値になります．

正解コード1

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

int main()
{
    //K番目までの入力
    int N, K;
    cin >> N >> K;
    priority_queue<int, vector<int>, greater<int>> pq;
    for(int i = 0; i < K; i++){
        int p;
        cin >> p;
        pq.push(p);
    }

    //1番最初の出力
    cout << pq.top() << endl;

    for(int i = K; i < N; i++){
        //K番目以降の入力
        int p;
        cin >> p;
        int q = pq.top();

        //pqと入力を比べる
        if(q < p){
            pq.pop();
            pq.push(p);
            cout << pq.top() << endl;
        }else{
            cout << q << endl;
        }
    }
}

実装2 ~setと二分探索~

　$K$番目までsetにinsertします．その時に1番小さい値をansとして保持します．
　1番最初の出力はansを出力します．
　それ以降は入力とansを比較し，ansが小さい場合はsetに入力をinsetし，ansをeraseします. そしたら**setからansより大きい1番近い値をを二分探索(upper_bound)**をしてその値が出力になります．
ansが大きい場合，ansを出力します．

正解コード2

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

int main()
{
    //入力
    int N, K;
    cin >> N >> K;
    set<int> st;
    int ans = N;

    for (int i = 0; i < K; i++)
    {
        int p;
        cin >> p;
        st.insert(p);
        if(p < ans) ans = p;
    }

    //1番最初の出力
    cout << ans << endl;

    for (int i = K; i < N; i++)
    {
        //それ以降の入力
        int p;
        cin >> p;

        // ansと入力を比べる
        if(p > ans){
            st.insert(p);
            st.erase(ans);
            auto iter = st.upper_bound(ans);
            ans = *iter;
            cout << ans << endl;
        }
        else{
            cout << ans << endl;
        }  
    }
}