K-th Number - shifth’s blog

SPOJ.com - Problem MKTHNUM
2104 -- K-th Number

最近 yukicoder で平方分割をやったんで(No.96 圏外です。 - shifth’s blog)蟻本にのっていた問題を復習してみた。
定数倍がきびしすぎるのか vector をやめて global 配列にして、 cin/cout を scanf/printf にして PKU でやっと通った。
SPOJ のほうはそもそも制限時間 0.6 秒くらいなんで平方分割じゃ無理なんだろう。
[l,r) の半区間を考えるのが実装のポイントかも。

#include <iostream>
#include <complex>
#include <sstream>
#include <string>
#include <algorithm>
#include <deque>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <vector>
#include <set>
#include <limits>
#include <cstdio>
#include <cctype>
#include <cmath>
#include <cstring>
#include <cstdlib>
#include <ctime>

typedef long long ll;
typedef unsigned long long ull;

#define FOR(i,a,b) for(int (i)=(a);i<(b);i++)
#define REP(i,n) FOR(i,0,n)
#define RANGE(vec) (vec).begin(),(vec).end()

using namespace std;

const int maxN = (int)(1E+5);
int a[maxN+1];
int sorted_a[maxN+1];
const int B = 1000;
vector<int> bucket[maxN/B+1];

// 愚直な実装だと O(m*n*log(n)) <= 10^5*5*10^3*16 = 8000000000 で間に合わない
// bucket法・segment 木などで高速化してやる。

// 平方分割
// 定数倍がかなり厳しく、vector をやめて配列に,cin/cout をやめて
// サーバが重くないときに PKU でやってやっと Accept する。SPOJ では何度やっても TLE
int main(int argc, char *argv[]) {
        int n,m;
        scanf("%d %d", &n, &m);
        REP(i,n)
        {
            scanf("%d", &a[i]);
            sorted_a[i] = a[i];
            bucket[i/B].push_back(a[i]);
        }
        //
        // まず x が区間で k 番目なら
        //  * x 未満の数値が k 個未満
        //  * x 以下の数値が k 個以上
        // なので区間がソート済みなら x に対する二分探索で k 番目の数が求まる。
        //
        // そこで数列を B サイズのバケットに分割しておき、各バケットでソートしておく
        //
        //  * 区間に完全に含まれるバケットなら二分探索
        //  * 区間がバケットの外に出てしまうものは個々にしらべる
        //
        // 各バケットでの計算量は O((n/B)log(B)+B) 程度
        //
        // B = √(n*log(n)) くらいにしておくと
        // O((n/B)*log(B)+B) = O(√(n*log(n)))
        //
        // 全体では前処理含めて O(n*log(n)+ m*√n*log^3/2(n)) 程度 (<=22000000)
        //
        sort(sorted_a,sorted_a+n);
        REP(i,n/B) // 末尾の B 未満のバケットは二分探索対象にならないのでそのままでよい
            sort(RANGE(bucket[i]));

        // 各クエリを処理
        REP(_, m)
        {
            int l,r,k;
            scanf("%d %d %d", &l,&r,&k);
            --l; // 実装を楽にするため [l,r) で計算する

            // a の値のうちどれかが解なので sorted_a のインデックスで二分探索をやる。
            // O(log(n)*((n/b)*log(b) + b))
            int ub = n-1;
            int lb = -1; // mid で 0 も選ばれるようにしておく
            while (lb+1 < ub)
            {
                int mid = (lb+ub)/2;
                int x   = sorted_a[mid];
                int tl = l;
                int tr = r;
                int cnt = 0; // [l,r) 内部の x 以下の数字の数

                // [l,r) 区間内の x 未満の個数を数える
                // bucket の外分
                // 左から
                while (tl < tr && tl % B) // バケットの先頭インデックスまで tl を動かす
                {
                    if (a[tl++] <= x)
                        ++cnt;
                }
                // 右から [l,r) の半開区間にするのはここの判定を楽にするため
                while (tl < tr && tr % B)
                {
                    // 半階区間なので tr-- でなく --tr
                    if (a[--tr] <= x)
                        ++cnt;
                }
                // 各 bucket ごとにやる
                while (tl < tr)
                {
                    int b = tl/B;
                    // バケット内の x 以下の数値を加算
                    cnt += upper_bound(RANGE(bucket[b]), x) - bucket[b].begin();
                    tl += B;
                }
                if (cnt >= k)
                    ub = mid;
                else
                    lb = mid;
            }
            printf("%d\n", sorted_a[ub]);
        }

        return 0;
}