AtCoder Regular Contest 101 D - Median of Medians

• 問題

D - Median of Medians

• 解法

であり、数列の長さはなので数列を実際につくるとTLEになるので別の方法を考えないといけない。
editorialがめちゃめちゃ分かりやすいので特に書くことなし。
BITでの反転数の求め方はここ
www.geeksforgeeks.org
が分かりやすいと思いました。

#include <iostream>
#include <vector>

using namespace std;

typedef long long int ll;

#define all(x) x.begin(),x.end()

struct BIT{
    ll N;
    vector<ll> bit;
    BIT(ll N) : bit(N+1),N(N){};

    ll sum(ll i){
        ll s = 0;
        while(i>0){
           s += bit[i];
           i -= i& -i;
        }
        return s;
    }

    void add(ll i,ll x){
        while(i <= N){
            bit[i] += x;
            i += i & -i;
        }
    }
};

//数列の値の範囲を[1,n]に変える
vector<ll> convert_array(vector<ll> z){
    ll n = z.size();
    vector<ll> temp = z;
    sort(all(temp));
    for(int i = 0; i < n; i++){
        z[i] = lower_bound(all(temp), z[i]) - temp.begin() + 1;
    }
    return z;
}

//反転数を求める
ll inverse_number(vector<ll> z){
    ll n = z.size();
    ll inv_num = 0;
    z = convert_array(z);

    BIT b(n);
    for(int i = 0; i <= n; i++){
        b.bit[i] = 0;
    }

    for(int i = 0; i < n; i++){
        inv_num += i - b.sum(z[i]);
        b.add(z[i],1);
    }

    return inv_num;
}

//数列の変形をしてxが中央値の候補かどうか調べる
bool check(ll x,vector<ll> z){
    ll n = z.size();
    vector<ll> b(n);
    for(int i = 0; i < n; i++){
        if(z[i] >= x) b[i] = 1;
        else b[i] = -1;
    }
    vector<ll> s(n+1,0);  //ここ長さn+1になることに注意
    for(int i = 1; i <= n; i++){
        s[i] = s[i-1] + b[i-1];
    }

    return (n*(n+1)/2 - inverse_number(s)) >= (n*(n+1)/2+1)/2;
}

int main()
{
    ll n;
    cin >> n;
    vector<ll> a(n);
    for(int i = 0; i < n; i++){
        cin >> a[i];
    }

    ll low = *min_element(all(a))-1, high = *max_element(all(a))+1;

    while(high - low > 1){
        ll mid = (low+high)/2;
        if(check(mid,a)) low = mid;
        else high = mid;
    }

    cout << low << endl;

    return 0;
}