備忘録

まとめておきたいことのメモ 主にプロコンのこと

ARC 043 B

今回はこの問題です。

arc043.contest.atcoder.jp

n個の問題があります。i番目の問題の難易度はd[i]です。
n個のうち、次の条件を満たすように問題を4つ選びます。

  • i + 1番目の問題の難易度はi問目の問題の難易度の2倍以上になる。(i = 1, 2, 3)

このとき、問題の選び方は何通りあるでしょうか、という問題です。数が大きくなるらしいので10^9 + 7でわったあまりを求めろ、とのことです。

わからなかったので解説を見ました。
1,2問目と3,4問目に分けて考えるのがキモ(っぽい)。

    #include <cstdio>
    #include <iostream>
    #include <cmath>
    #include <cstring> 
    #include <sstream>
    #include <iostream>
    #include <algorithm>
    #include <cstdlib>
    #include <map>
    #include <queue>
    #include <utility>
    #include <vector>
    #include <set>
    #include <memory.h>
    #include <iomanip>
    #include <bitset>
    #include <list>
    #include <stack>
     
    using namespace std;
     
    #define pai 3.14159265359
    #define mod 1000000007
     
    int main()
    {
        int n;
        int d[100001];
        long long int a[100001] = {}, b[100001] = {};
        cin >> n;
        for(int i = 1; i <= n; i++){
            cin >> d[i];
        }
        sort(d + 1, d + n + 1);
        // a[i] : 難易度が2 * d[i]以上となる問題の個数 つまり3番目の問題の難易度をd[i]としたとき4番目として選べる問題の数
        for(int i = 1; i <= n; i++){
            int left = i, right = n;
            if(d[i] * 2 > d[n]) continue;
            while(right - left > 1){
                int mid = (left + right) / 2;
                if(d[mid] >= 2 * d[i]) right = mid;
                else left = mid;
            }
            a[i] = n - left;
            // cout << i << " " << a[i] << endl;
        }
        // b[i] : 難易度がd[i] / 2以下となる問題の個数 つまり2番目の問題の難易度をd[i]としたとき1番目のとして選べる問題の数
        for(int i = 1; i <= n; i++){
            int left = 1, right = i;
            if(d[i] / 2 < d[1]) continue;
            while(right - left > 1){
                int mid = (left + right) / 2;
                if(2 * d[mid] <= d[i]) left = mid;
                else right = mid;
            }
            b[i] = right - 1;
            // cout << i << " " << b[i] << endl;
        }
        // ここからaの値を更新 
        // a[i] : 3番目の問題の難易度がd[i]以上のときの3問目と4問目の組み合わせの数
        for(int i = n - 1; i >= 1; i--){
            a[i] += a[i + 1];
            a[i] = a[i] % mod;
        }
        int cnt = 1;
        long long int ans = 0;
        for(int i = 1; i <= n; i++){
            // 2番目の問題の難易度がd[i]であるとする
            while(cnt <= n && d[cnt] < 2 * d[i]) cnt++;
            // cout << d[i] << " " << d[cnt] << endl;
            // cout << b[i] << " " << a[cnt] << endl;
            ans += (b[i] * a[cnt]) % mod;
            ans = ans % mod;
        }
        cout << ans << endl;
        return 0;
    }

コメント通り、まず
a[i] : 難易度が2 * d[i]以上となる問題の個数 つまり3番目の問題の難易度をd[i]としたとき4番目として選べる問題の数
b[i] : 難易度がd[i] / 2以下となる問題の個数 つまり2番目の問題の難易度をd[i]としたとき1番目のとして選べる問題の数
としてそれぞれ二分探索で求めます。
さらに、
a[i] : 3番目の問題の難易度がd[i]以上のときの3問目と4問目の組み合わせの数
に更新します。後ろから順番に足せばオッケーです。
あとは、2番目の問題の難易度をd[i]である場合を順番に足していきます。d[cnt] >= 2 * d[i]となるようにcntを求めています。3問目の難易度がd[cnt]以上のときの3問目と4問目の組み合わせの数 * 2問目の難易度がd[i]で1問目の難易度がd[i] / 2以下になるような組み合わせの数を計算すればオッケーです。