poj2104 K-th Number (线段树实现)

原题链接：K-th Number

题意：给定一个数列a1,a2,...,an和m个查询操作。对于每个查询(i,j,k)，输出ai,...,aj的升序排列后的第k个数

//用线段树实现 
#include <cstdio>
#include <vector>
#include <algorithm> 
using namespace std;
const int ST_SIZE = (1 << 18) - 1;
const int MAX_N = 100000 + 5;
const int MAX_M = 5000 + 5;
//输入 
int N, M;
int A[MAX_N];
int I[MAX_M], J[MAX_M], K[MAX_M];

int nums[MAX_N]; //对A排序之后的结果 
vector<int > dat[ST_SIZE];	//线段树数据 

//k是节点的编号，和区间[l, r)对应 
void init(int k, int l, int r){	//构建线段树 
	if(r - l == 1){	//叶节点 
		dat[k].push_back(A[l]);
	}else {	//父节点 
		int lch = k * 2 + 1, rch = k * 2 + 2;
		init(lch, l, (l + r) / 2);
		init(rch, (l + r) / 2, r);
		dat[k].resize(r - l);
		//利用STL的merge函数把两个儿子的数列合并
		merge(dat[lch].begin(), dat[lch].end(), 
			  dat[rch].begin(), dat[rch].end(), 
			  dat[k].begin()); 
	}
}

//[i, j)中不超过x的数的个数
//k是节点的编号，和区间[l, r)对应 
int query(int i, int j, int x, int k, int l, int r){
	if(j <= l || r <= i){	//不相交 
		return 0;
	}else if(i <= l && r <= j){	//完全包含 
		return upper_bound(dat[k].begin(), dat[k].end(), x) - dat[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;
	}
} 

void solve(){
	for(int i = 0;i < N;i ++){
		nums[i] = A[i];
	}
	sort(nums, nums + N);
	init(0, 0, N);
	for(int i = 0;i < M;i ++){	//查找[l, r)中的第k个数 
		int l = I[i], r = J[i] + 1, k = K[i];
		int lb = -1, ub = N - 1;
		while(ub - lb > 1){
			int mid = (ub + lb) / 2;
			int c = query(l, r, nums[mid], 0, 0, N);
			if(c >= k){
				ub = mid;
			} else {
				lb = mid;
			}
		}
		printf("%d\n", nums[ub]);
	}
}

int main(){
	int T, x;
	scanf("%d%d", &N, &M);
	for(int i = 0;i < N;i ++){
		scanf("%d", &A[i]);
	}
	for(int i = 0;i < M;i ++) {
		scanf("%d%d%d", &I[i], &J[i], &K[i]);
		I[i] --; J[i]--;	//下标是从0开始的 
	}
	solve();
	return 0;
} 

//用平方分割超时代码： 
#include <cstdio>
#include <vector>
#include <algorithm> 
using namespace std;
const int B = 1000;	//桶的大小
const int MAX_N = 100000;
const int MAX_M = 5000;
//输入 
int N, M;
int A[MAX_N];
int I[MAX_M], J[MAX_M], K[MAX_M];

int nums[MAX_N]; //对A排序之后的结果 
vector<int > bucket[MAX_N / B];	//每个桶排序之后的结果 

void solve(){
	for(int i = 0;i < N;i ++){
		bucket[i / B].push_back(A[i]);
		nums[i] = A[i];
	}
	sort(nums, nums + N);
	//虽然每B个一组剩下的部分所在的桶没有排序，但是不会产生问题
	for(int i = 0;i < N / B;i ++){
		sort(bucket[i].begin(), bucket[i].end());
	}
	for(int i = 0;i < M;i ++){	//求[l, r)区间中第k个数 
		int l = I[i], r = J[i] + 1, k = K[i];
		int lb = -1, ub = N - 1;
		while(ub - lb > 1){	//二分查找 
			int md = (lb + ub) / 2;
			int x = nums[md];
			int tl = l, tr = r, c = 0;
			//区间两端多出的部分
			while(tl < tr && tl % B != 0){
				if(A[tl ++] <= x){
					c ++;
				}
			} 
			while (tl < tr && tr % B != 0){
				if(A[-- tr] <= x){
					c ++;
				}
			} 
			//对每一个桶进行计算
			while(tl < tr){
				int b = tl / B;
				c += upper_bound(bucket[b].begin(), bucket[b].end(), x)
					 - bucket[b].begin();
				tl += B;
			} 
			if(c >= k){
				ub = md;
			}else{
				lb = md;
			}
		}
		printf("%d\n", nums[ub]);
	}
}

int main(){
	int T, x;
	scanf("%d%d", &N, &M);
	for(int i = 0;i < N;i ++){
		scanf("%d", &A[i]);
	}
	for(int i = 0;i < M;i ++){
		scanf("%d%d%d", &I[i], &J[i], &K[i]);
		I[i] --; J[i] --;
	}
	solve();
	return 0;
}