K-th Number (Segment Tree)

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

• segment tree による手法だと一発で AC した。(SPOJ)
• メモリをたくさん使う分やはり速い。
• merge は std::back_inserter を使ってもよいが事前に resize しておくほうが速いはず。
• バケット法はこちらK-th Number - shifth’s blog

class KthNumber {
public:
    // セグメント木による手法
    // 各ノードに数列をもたせる。
    struct SegTree {
        vector<vector<int>> data;
        SegTree(const vector<int> &a) {
                int k = 1;
                // 一番下の段が a を含むようにする
                while (k < (int)a.size()) k <<= 1;
                data.resize(2*k); // 2*(最下段の長さ) が SegTree の必要メモリサイズ
                init(a, 0, 0, k); // 一番上の段は [0,k) に対応
        }
        // <--------- log(K) ----------->
        // 1*K + 2*(K/2) + 4*(K/4) + ...  = O(K*log(K))
        // memory used : O(K*log(K))
        //
        // data[k] ノードが [l,r) を表現する
        void init(const vector<int> &a, int k, int l, int r) {
                if (r-l == 1) // 葉
                {
                    if (l < (int)a.size())
                        data[k].push_back(a[l]);
                }
                else
                {
                    int lch = k*2+1;
                    int rch = k*2+2;
                    init(a, lch, l, (l+r)/2);
                    init(a, rch, (l+r)/2, r);
                    data[k].resize(data[lch].size()+data[rch].size()); // merge 前にメモリ確保
                    merge(RANGE(data[lch]), RANGE(data[rch]), data[k].begin());
                }
        }
        //  O(log^2(K))
        // [i,j) 内の x の数を数える
        int query(int i, int j, int x,   // query 情報
                  int k, int l, int r) { // node の情報
                if (j <= l || r <= i)
                    return 0;
                else if (i <= l && r <= j) // query 区間がノード管理区間にすっぽり含まれるとき
                    return upper_bound(RANGE(data[k]), x)-data[k].begin();
                else
                {
                    int lc = query(i, j, x, k*2+1, l, (l+r)/2);
                    int rc = query(i, j, x, k*2+2, (l+r)/2, r);
                    return lc+rc;
                }
        }
        int query(int i, int j, int x) { return query(i, j, x, 0, 0, data.size()/2); }
    };
    // O(n*log(n)+m*log^3(n))
    void solve_segtree() {
            int n,m;
            cin>>n>>m;
            vector<int> a(n);
            vector<int> sorted_a(n); // sort 済み a
            REP(i,n)
            {
                cin>>a[i];
                sorted_a[i] = a[i];
            }
            sort(RANGE(sorted_a));
            // sort 済み配列から segtree を構築
            SegTree sgt(a);

            // O(m*log^3(n))
            REP(_,m)
            {
                int l,r,k;
                cin>>l>>r>>k;
                --l;
                int lb = -1;
                int ub = a.size()-1;

                // binary search
                while (ub-lb > 1)
                {
                    int md = (lb+ub)/2;
                    int x = sorted_a[md];
                    int c = sgt.query(l,r,x);
                    if (c >= k)
                        ub = md;
                    else
                        lb = md;
                }
                cout<<sorted_a[ub]<<endl;
            }
    }
    void solve() {
            // solve_bucket();
            solve_segtree();
    }
 };