Codeforces Round #326 Duff in the Army（主席树+LCA）

题意：给出一棵树，每个节点有一个权值集合，现在有q个询问，query（u,v,a）表示询问u到v之间的前a小的数。

思路：主席树+LCA。

关于这道题的简化版详见http://blog.csdn.net/u014664226/article/details/47850199

这道题不同之处在于每个节点有多个权值，而且要输出前a小的数，只需要在建树时和查询时稍微改动一下即可。

查询时从第1小的查到第k小的。

#include<cstdio>
#include<cstring>
#include<cmath>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#include<vector>
#include<map>
#include<queue>
#include<stack>
#include<string>
#include<map>
#include<set>
#include<ctime>
#define eps 1e-6
#define LL long long
#define pii (pair<int, int>)
//#pragma comment(linker, "/STACK:1024000000,1024000000")
using namespace std;

//const int maxn = 100000 + 100;
//const int INF = 0x3f3f3f3f;

const int maxn = 100000+10000;
const int M = 2000000;
int n, q, m, tot;
int t[maxn], fa[maxn];
int T[maxn], lson[M], rson[M], c[M];
vector<int> C[maxn];
struct Quest {
int l, r, k;
} quest[maxn];

int build(int l, int r) {
int root = tot++;
c[root] = 0;
if(l != r) {
int mid = (l+r) >> 1;
lson[root] = build(l, mid);
rson[root] = build(mid+1, r);
}
return root;
}

int Insert(int root, int pos, int val) {
int newroot = tot++, tmp = newroot;
int l = 1, r = m;
c[newroot] = c[root] + val;
while(l < r) {
int mid = (l+r)>>1;
if(pos <= mid) {
lson[newroot] = tot++; rson[newroot] = rson[root];
newroot = lson[newroot]; root = lson[root];
r = mid;
}
else {
rson[newroot] = tot++; lson[newroot] = lson[root];
newroot = rson[newroot]; root = rson[root];
l = mid+1;
}
c[newroot] = c[root] + val;
}
return tmp;
}

int Query(int l_root, int r_root, int lca, int k) {
int l = 1, r = m, lca_root = T[lca], fa_root = fa[lca];
bool tag = 0;
int tmp = c[l_root]+c[r_root]-c[lca_root]-c[fa_root];
if(tmp < k) return -1;
while(l < r) {
int mid = (l+r) >> 1;
int tmp = c[lson[l_root]]+c[lson[r_root]]-c[lson[lca_root]]-c[lson[fa_root]];
if(tmp >= k) {
r = mid;
l_root = lson[l_root]; r_root = lson[r_root]; lca_root = lson[lca_root]; fa_root = lson[fa_root];
tag = 1;
}
else {
l = mid + 1;
k -= tmp;
l_root = rson[l_root]; r_root = rson[r_root]; lca_root = rson[lca_root]; fa_root = rson[fa_root];
}
}
return l;
}

int pnt[maxn], lca[maxn];
bool vis[maxn];
vector<int> G[maxn], query[maxn], num[maxn];
int find(int x) {
if(x == pnt[x]) return x;
return pnt[x] = find(pnt[x]);
}
void dfs_lca(int u) {
vis[u] = 1; pnt[u] = u;
int sz1 = G[u].size();
for(int i = 0; i < sz1; i++) {
int v = G[u][i];
if(vis[v]) continue;
fa[v] = T[u];
dfs_lca(v);
pnt[v] = u;
}
int sz2 = query[u].size();
for(int i = 0; i < sz2; i++) {
int v = query[u][i];
if(vis[v]) lca[num[u][i]] = find(v);
}
}

void dfs_ZXTree(int cur, int fa) {
int sz = G[cur].size();
for(int i = 0; i < sz; i++) {
int u = G[cur][i];
if(u == fa) continue;
T[u] = T[cur];
for(int j = 0; j < C[u].size(); j++) T[u] = Insert(T[u], C[u][j], 1);
dfs_ZXTree(u, cur);
}
}

int main() {
// freopen("input.txt", "r", stdin);
cin >> n >> m >> q;
for(int i = 1; i < n; i++) {
int u, v; scanf("%d%d", &u, &v);
G[u].push_back(v);
G[v].push_back(u);
}
for(int i = 1, u; i <= m; i++) scanf("%d", &u), C[u].push_back(i);

build(1, m);

T[1] = T[0];
for(int i = 0; i < C[1].size(); i++) T[1] = Insert(T[1], C[1][i], 1);

dfs_ZXTree(1, -1);

for(int i = 0; i < q; i++) {
int l, r, k; scanf("%d%d%d", &l, &r, &k);
quest[i].l = l; quest[i].r = r; quest[i].k = k;
query[l].push_back(r); query[r].push_back(l);
num[l].push_back(i); num[r].push_back(i);
}
fa[1] = T[0];
dfs_lca(1);
for(int i = 0; i < q; i++) {
int cnt = 0;
vector<int> ans;
for(int j = 1; j <= quest[i].k; j++) {
int tmp = Query(T[quest[i].l], T[quest[i].r], lca[i], j);
if(tmp > 0) ans.push_back(tmp);
else break;
}
int sz = ans.size();
printf("%d", sz);
for(int j = 0; j < sz; j++) printf(" %d", ans[j]);
puts("");
}
return 0;
}