[C++][AtCoder/競プロ]二重ループ使用の際に、AtCoderにてTLE(実行時間制限超過)となるのを防ぎたい。

Question

前提・実現したいこと

こんにちは。
現在、こちらの問題を解いています。
そこで、実行時間制限超過のエラーが出てしまっているので、直したいと考えています。

ソースコード

こちらの問題を解いていて、コードを以下のように書きました。

#include <bits/stdc++.h>

using namespace std;

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    vector<int> x(Q);

    int count=0;

    for (int i=0; i<N; i++) {
        cin >> a.at(i);
    }
    for (int i=0; i<Q; i++) {
        cin >> x.at(i);
    }

    sort(a.begin(), a.end());
    for (int i=0; i<Q; i++) {
        for (int j=0; j<N; j++) {
            if (a.at(j) >= x.at(i)) {
                count++;
            }
            if (j == N-1) {
                cout << count << endl;
                count = 0;
            }
        }
    }
}

発生している問題

上記のコードを提出すると、TLE(実行時間制限超過)と表示されてしまいます。
恐らく、二重ループを使用しているところで計算量が大きくなってしまっているため、エラーが発生しているのではないかと思っております。
しかしながら、二重ループ以外で、どう書き換えればいいのかが分かりません。

ざっくりとした質問になってしまいましたが、ご教授いただけると幸いです。

よろしくお願い致します。

Comments

[Crimson_Tide]

二重ループで最大で10^9 × 10^9　してるので現実的ではないですね。
線形探索では厳しいと思います。

https://atcoder.jp/contests/abc231/editorial/3074
公式解説が既にあるのでご覧になるとよいかと思います。
ソートした配列に対して二分探索することでｘより大きくなる要素番号を見つけ、Nからその要素番号を引いてますね。

https://atcoder.jp/contests/abc231/editorial/3075
さらにもう１つの解説動画でクエリ(xの配列)のソートをしてます。
具体的な処理は見ていませんが、恐らく二分探索の範囲を徐々に狭め探索回数の軽減につなげているではないかと思います。

[Kaguya_4869]

Crimson_Tide様
公式解説の更に詳しい解説、ありがとうございます。
二分探索が少し苦手なため、今まで避けてきていたのですが、
この機会に正面から向き合ってみようと思います！
ありがとうございます。

[episteme]

> https://atcoder.jp/contests/abc231/editorial/3075
> さらにもう１つの解説動画でクエリ(xの配列)のソートをしてます。
> 具体的な処理は見ていませんが、恐らく二分探索の範囲を徐々に狭め探索回数の軽減につなげているで
> はないかと思います。

ホントにこっちの方が速いか疑わしい...二分検索してないっポいし。

Answer 1

配列aをソートしたならば、全ての要素を探索せずとも、x以上である最小の要素a[i]を見つければ、身長がx以上の生徒の人数は(N-i)人と求められます。また、最小の要素が見つからなかった場合も、i=Nとすることで同様に求められます。

そして、最小の要素a[i]を求めるには、線形探索を使うと一回の質問ごとにO(N)かかってしまうので、二分探索を使います。すると質問ごとの計算量をO(logN)に抑えることができ、全体の計算量がO(NlogN)となって通ります。なお、二分探索を使うには配列aがソートされている必要があります。

二分探索にはlower_bound関数を使いましょう。ただし、lower_boundの返り値はイテレーターなので、返り値からa.begin()を引くことで、iを求めることができます。(今回はa.end()から返り値を引けば、生徒の人数を直接求められます。)

コード例:

#include <bits/stdc++.h>

using namespace std;

int main()
{
    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    vector<int> x(Q);

    for (int i = 0; i < N; i++)
    {
        cin >> a.at(i);
    }
    for (int i = 0; i < Q; i++)
    {
        cin >> x.at(i);
    }

    sort(a.begin(), a.end());
    for (int i = 0; i < Q; i++)
    {
        cout << a.end() - lower_bound(a.begin(), a.end(), x[i]) << endl;
        //cout << N - (lower_bound(a.begin(), a.end(), x[i]) - a.begin()) << endl;
    }
}

二分探索について詳しく知りたい場合は、こちらのサイトが参考になると思います。
https://pyteyon.hatenablog.com/entry/2019/02/20/194140

Comments

[episteme]

> a.end() - lower_bound(a.begin(), a.end(), x[i])

あるいは distance(lower_bound(a.begin(), a.end(), x[i]), a.end())

[Kaguya_4869]

luuguas様

回答ありがとうございます。
二分探索についての分かりやすいサイトまで紹介してくださり、
本当にありがとうございます。
理解することができました！
これから、もう少し二分探索についての問題等を解いていきたいと思います。

[luuguas]

二分探索は競プロの典型テクニックなので、使いこなせるようにしておきましょう。(自戒)

Answer 2

vector<int> a(N); をソートしたのなら、逐次比較しつつ勘定する必要はないのでは?
※ なんのためにソートしたんです?

a が昇順にソートされているなら、
任意の身長xに対し x <= a[t] なる最小の要素:a[t]を見付ければ、
そこから aの末尾まで の距離(人数)が求める答ですよね?

[追記] やってみた:

#include <iostream>
#include <algorithm>
#include <vector>
#include <iterator>

int main() {
    using namespace std;

    int N, Q;
    cin >> N >> Q;

    vector<int> a(N);
    vector<int> x(Q);

    for ( int& item : a ) cin >> item;
    for ( int& item : x ) cin >> item;

    sort(a.begin(), a.end());
    for (int i=0; i<Q; i++) {
      cout << distance(lower_bound(a.begin(), a.end(), x.at(i)), a.end()) << endl;
    }
}

Comments

[Kaguya_4869]

episteme様

昨日の多次元配列の質問に続き、こちらにも回答してくださり、
ありがとうございます。
確かに、私のコードですと、ソートしたにも関わらず、
それを活かせていませんでした。
上の方々のコメントや回答から、
x<=a[t]になる最小の要素a[t]を見つけるには、
二分探索が必要そうだということでしたので、
それを使ってみようと思います。
ありがとうございます。

[episteme]

追記しました。