【bzoj4771】七彩树 树链的并+STL-set+DFS序+可持久化线段树
题目描述
给定一棵n个点的有根树，编号依次为1到n，其中1号点是根节点。每个节点都被染上了某一种颜色，其中第i个节点的颜色为c[i]。如果c[i]=c[j]，那么我们认为点i和点j拥有相同的颜色。定义depth[i]为i节点与根节点的距离，为了方便起见，你可以认为树上相邻的两个点之间的距离为1。站在这棵色彩斑斓的树前面，你将面临m个问题。每个问题包含两个整数x和d，表示询问x子树里且depth不超过depth[x]+d的所有点中出现了多少种本质不同的颜色。请写一个程序，快速回答这些询问。
输入
第一行包含一个正整数T(1<=T<=500)，表示测试数据的组数。
每组数据中，第一行包含两个正整数n(1<=n<=100000)和m(1<=m<=100000)，表示节点数和询问数。
第二行包含n个正整数，其中第i个数为c[i](1<=c[i]<=n)，分别表示每个节点的颜色。
第三行包含n-1个正整数，其中第i个数为f[i+1](1<=f[i]<i)，表示节点i+1的父亲节点的编号。
接下来m行，每行两个整数x(1<=x<=n)和d(0<=d<n)，依次表示每个询问。
输入数据经过了加密，对于每个询问，如果你读入了x和d，那么真实的x和d分别是x xor last和d xor last，
其中last表示这组数据中上一次询问的答案，如果这是当前数据的第一组询问，那么last=0。
输入数据保证n和m的总和不超过500000。
输出
对于每个询问输出一行一个整数，即答案。
样例输入
1
5 8
1 3 3 2 2
1 1 3 3
1 0
0 0
3 0
1 3
2 1
2 0
6 2
4 1
样例输出
1
2
3
1
1
2
1
1
题解
树链的并+STL-set+DFS序+可持久化线段树
如果没有深度限制，那么对于一种颜色，子树内包含该颜色的节点为：所有该颜色节点到根节点路径覆盖的所有节点，即树链的并。
因此对于每一种颜色求树链的并，支持链加操作；查询单个点的时候就是查询单点值。树上差分后转变为单点加、子树求和，使用DFS序转化为区间问题后使用数据结构维护。
那么有深度限制呢？我们按照深度维护可持久化线段树，第 $i$ 个版本我们只考虑深度小于等于 $i$ 的节点的影响。此时需要使用STL-set维护每个颜色的树链的并。
查询时直接查询depth[x]+d版本对应的可持久化线段树中，x节点子树内的权值和即可。
时间复杂度 $O(n\log n)$ 。
#include <set>
#include <cstdio>
#include <cstring>
#include <algorithm>
#define N 100010
using namespace std;
set<int> s[N];
set<int>::iterator it;
int c[N] , head[N] , to[N] , next[N] , cnt , fa[N][20] , deep[N] , log[N] , pos[N] , ref[N] , last[N] , tp;
int id[N] , sum[N << 6] , ls[N << 6] , rs[N << 6] , root[N] , tc;
inline void add(int x , int y)
{
to[++cnt] = y , next[cnt] = head[x] , head[x] = cnt;
}
void dfs(int x)
{
int i;
pos[x] = ++tp , ref[tp] = x;
for(i = 1 ; (1 << i) <= deep[x] ; i ++ ) fa[x][i] = fa[fa[x][i - 1]][i - 1];
for(i = head[x] ; i ; i = next[i]) fa[to[i]][0] = x , deep[to[i]] = deep[x] + 1 , dfs(to[i]);
last[x] = tp;
}
inline int lca(int x , int y)
{
int i;
if(deep[x] < deep[y]) swap(x , y);
for(i = log[deep[x] - deep[y]] ; ~i ; i -- )
if(deep[x] - deep[y] >= (1 << i))
x = fa[x][i];
if(x == y) return x;
for(i = log[deep[x]] ; ~i ; i -- )
if(deep[x] >= (1 << i) && fa[x][i] != fa[y][i])
x = fa[x][i] , y = fa[y][i];
return fa[x][0];
}
bool cmp(int a , int b)
{
return deep[a] < deep[b];
}
void update(int p , int a , int l , int r , int x , int &y)
{
y = ++tc , sum[y] = sum[x] + a;
if(l == r) return;
int mid = (l + r) >> 1;
if(p <= mid) rs[y] = rs[x] , update(p , a , l , mid , ls[x] , ls[y]);
else ls[y] = ls[x] , update(p , a , mid + 1 , r , rs[x] , rs[y]);
}
int query(int b , int e , int l , int r , int x)
{
if(b <= l && r <= e) return sum[x];
int mid = (l + r) >> 1 , ans = 0;
if(b <= mid) ans += query(b , e , l , mid , ls[x]);
if(e > mid) ans += query(b , e , mid + 1 , r , rs[x]);
return ans;
}
int main()
{
int T;
scanf("%d" , &T);
while(T -- )
{
cnt = tp = tc = 0;
int n , m , i , p = 1 , x , y , lastans = 0;
scanf("%d%d" , &n , &m);
for(i = 1 ; i <= n ; i ++ ) scanf("%d" , &c[i]) , s[i].clear() , id[i] = i , head[i] = 0;
for(i = 2 ; i <= n ; i ++ ) scanf("%d" , &x) , add(x , i) , log[i] = log[i >> 1] + 1;
dfs(1) , sort(id + 1 , id + n + 1 , cmp);
for(root[0] = i = 0 ; i < n ; i ++ )
{
if(i) root[i] = root[i - 1];
while(p <= n && deep[id[p]] <= i)
{
x = y = 0;
it = s[c[id[p]]].lower_bound(pos[id[p]]);
if(it != s[c[id[p]]].end()) y = ref[*it];
if(it != s[c[id[p]]].begin()) x = ref[*--it];
update(pos[id[p]] , 1 , 1 , n , root[i] , root[i]);
if(x) update(pos[lca(x , id[p])] , -1 , 1 , n , root[i] , root[i]);
if(y) update(pos[lca(y , id[p])] , -1 , 1 , n , root[i] , root[i]);
if(x && y) update(pos[lca(x , y)] , 1 , 1 , n , root[i] , root[i]);
s[c[id[p]]].insert(pos[id[p]]) , p ++ ;
}
}
while(m -- )
{
scanf("%d%d" , &x , &y) , x ^= lastans , y = min(deep[x] + (y ^ lastans) , n - 1);
printf("%d\n" , lastans = query(pos[x] , last[x] , 1 , n , root[y]));
}
}
return 0;
}
