AtCoder ABC 077 C - Snuke Festival (ARC 084 C) (緑色, 300 点)

lower_bound の練習に！！！　あと、「3 つのものを考えるときは、真ん中を固定して考える」という考え方の典型。

• 問題へのリンク

問題概要

3 つの数列 (長さ )

• 
• 
• 

が与えられる。各数列から要素 を選ぶ方法のうち、

を満たすものの個数を求めよ。

制約

• 

考えたこと

単純にすべてのペア を考える方法では の計算量となりますので、工夫が必要です。

まずは考えやすくなるように、各数列は小さい順にソートされているものとしましょう (ソートは でできます)。さて、3 つ組について考えるときは、真ん中を固定して考えるのが定石です。つまり、 を固定して考えてみましょう。このとき、

• のうち、 未満の個数を数える ( とする)
• のうち、 以上の個数を数える ( とする)

というようにします (下図参照)。 を満たす組の個数は となります。これを について足すことで答えが求められます。

と を求める

これらはいずれも二分探索 (特に lower_bound) を使って求められます！

まず のうち、 未満の個数を数えてみましょう。lower_bound を使うと、次のようにして を満たす最小の が求められます。

int k = lower_bound(a.begin(), a.end(), b[j]) - a.begin();

そしてこのとき、 を満たす の個数は になります。こうして であることがわかりました。

についても同様に求められます。具体的には、 (不等号に等号が付くことに注意) を満たす の個数を求めて、それを から引けばよいでしょう。そのような の個数は、lower_bound() の代わりに upper_bound() を用いることで求められます。

計算量・コード

各 に対して、 と は で求められます。よって全体の計算量は となることがわかりました。

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int main() {
    // 入力
    int N;
    cin >> N;
    vector<long long> a(N), b(N), c(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    for (int i = 0; i < N; ++i) cin >> b[i];
    for (int i = 0; i < N; ++i) cin >> c[i];

    // ソートする
    sort(a.begin(), a.end());
    sort(b.begin(), b.end());
    sort(c.begin(), c.end());

    // b[j] を固定して考える
    long long res = 0;
    for (int j = 0; j < N; ++j) {
        long long Aj = lower_bound(a.begin(), a.end(), b[j]) - a.begin();
        long long Cj = N - (upper_bound(c.begin(), c.end(), b[j]) - c.begin());
        res += Aj * Cj;
    }
    cout << res << endl;
}