题目描述
给定一棵树,结点由 1
至 n
编号,其中结点 1
是树根。
树的每个点有一个颜色 Ci
。
如果一棵树中存在的每种颜色的结点个数都相同,则我们称它是一棵颜色平衡树。
求出这棵树中有多少个子树是颜色平衡树。
输入格式
输入的第一行包含一个整数 n
,表示树的结点数。
接下来 n
行,每行包含两个整数 Ci,Fi
,用一个空格分隔,表示第 i
个结点的颜色和父亲结点编号。
特别地,输入数据保证 F1
为 0
,也即 1
号点没有父亲结点。
保证输入数据是一棵树。
输出格式
输出一行包含一个整数表示答案。
数据范围
对于 30%
的评测用例,n≤200
,Ci≤200
;
对于 60%
的评测用例,n≤5000
,Ci≤5000
;
对于所有评测用例,1≤n≤200000
,1≤Ci≤200000
,0≤Fi<i
。
样例
输入样例:
6
2 0
2 1
1 2
3 3
3 4
1 4
输出样例:
4
样例解释:
编号为 1,3,5,6的4个结点对应的子树为颜色平衡树。
时间复杂度
O(n * logn)
参考文献
洛谷上某位大佬,链接如下:
以下只是我对大佬代码的理解,希望能给大家带来些许帮助。(渴望自己也变强的那一天,贵在坚持!)
https://www.luogu.com.cn/problem/solution/P9233
C++ 代码
//首先,从大佬给出的题解可以看出,就是要保存重儿子的对cnt[],ccnt[]的改变的贡献,但是并不保存轻儿子对这两个数组的状态的改变的贡献的。
//从所给样例中可以看出,子树也是包括了这个'根节点'的。
#include <bits/stdc++.h>
#define rep(x,y,z) for(int x=(y);x<=(z);x++)
// #define per(x,y,z) for(int x=(y);x>=(z);x--)
// #define debug(format...) fprintf(stderr, format)
// #define fileIO(s) do{freopen(s".in","r",stdin);freopen(s".out","w",stdout);}while(false)
// #define likely(exp) __builtin_expect(!!(exp), 1)
// #define unlikely(exp) __builtin_expect(!!(exp), 0)
using namespace std;
typedef long long ll;
/*
mt19937 rnd(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::system_clock::now().time_since_epoch()).count());
int randint(int L, int R) {
uniform_int_distribution<int> dist(L, R);
return dist(rnd);
}*/
template<typename T> void chkmin(T& x, T y) {if(x > y) x = y;} //果然,我就知道这里是有一些行代码是可有可无的。
template<typename T> void chkmax(T& x, T y) {if(x < y) x = y;}
const int N = 2e5+5;
int n, c[N], f[N], sz[N], son[N], cnt[N], ccnt[N], ans;
//n就是结点的个数的
//c[i]就是存储每个结点的颜色的
//f[i]就是存储第i个结点的父节点的。
//sz[i]就是存储以i为根节点的包括这个根的子树的结点的数量的。由于下面dfs()函数中最开始设置sz[u] = 1可以看出是要包括根结点的。
//son[i]就是存储以第i个点为根节点的重儿子的结点编号的。
//从下面的calc()函数中可以看得出,cnt[i]就是存储第i种颜色的出现的次数的。
//也可以总下面的calc()函数和大佬的题解中看出,ccnt[i]存储的就是存储颜色的出现次数的出现次数,比如1种颜色2个结点,是有几种颜色会是这样的。存储的就是这个的。
vector<int> e[N]; //这个是最开始的一行的。
//像我下面这样写就是不行的,下面这样写的话就是不行的,上面这一行就是可以视为是一个二维数组的,形如e[i][j],且第二维的大小是不定长度的。
// vector<int> e;
//e[i]就是用来存储第i个结点的所有儿子结点的,孙子结点是不会算进去的。
void dfs(int u) {
sz[u] = 1; //由于最开始设置sz[u] = 1可以看出是要包括根结点的。
for(int v : e[u]) {
dfs(v);
sz[u] += sz[v];
if(sz[v] > sz[son[u]]) son[u] = v;
}
}
void add(int u, int dt) {
--ccnt[cnt[c[u]]]; //这个就是说上一种出现次数减去1的,下一种出现次数加上1的。
cnt[c[u]] += dt;
++ccnt[cnt[c[u]]];
for(int v : e[u]) add(v, dt); //也就是说对于当前这个结点的所有子孙结点,都是会这样进行的。
}
void calc(int u, bool save) {
for(int v : e[u]) if(v != son[u]) calc(v, false); //如果是轻儿子的话就是设置save是false的。
if(son[u]) calc(son[u], true); //如果是重儿子的话就是设置save是true的。
--ccnt[cnt[c[u]]]; //这一行和下面两行是对应的,由于是当前这种颜色多了一个的,则上一个出现次数显然会减去1,下面的多一次的出现次数的出现次数显然也就是会加上1的。
++cnt[c[u]]; //表面上看起来上面这一行和下面这一行是一样的,但是实际上不是的,经过这一行之后cnt[c[u]]就是发生了改变的。
++ccnt[cnt[c[u]]]; //接着上面的2行的。
for(int v : e[u]) if(v != son[u]) add(v, 1); //这个就是不是重儿子才会进行的。
if(cnt[c[u]] * ccnt[cnt[c[u]]] == sz[u]) ++ans; //显然如果是颜色平衡树的话这个就是会满足的。
//对于上面这一行的解释,其实我第一遍确实是可以懂,但是以后未必可以看的懂,希望这些补充会对下次看到这个程序的我有所帮助。
//对于每一颗颜色平衡树,它的cnt[c[u]]是固定的,ccnt[cnt[c[u]]]也是固定的,只要等于sz[u]就是视为是一棵颜色平衡树的。
if(!save) add(u, -1); //这个就是说如果是save = false的情况就要进行减去的。而且是第u这个点,和u的所有子结点都给进行--的。
//就是说如果这个u点不是重儿子的点的话,视为是先使用再恢复现场的。
}
int main() {
scanf("%d", &n);
rep(i, 1, n) { //这个说实话是和我自己的习惯不太一样的,这个在最上面是进行了define的。
scanf("%d%d", &c[i], &f[i]);
if(f[i]) e[f[i]].push_back(i);
}
dfs(1); //这个就是在进行一次深度优先遍历的。这样的话每个结点的重儿子,以及每个子树的结点的数量都是在sz[i],son[i]数组中进行了保存的。
calc(1, true); //这个函数里面是灵活使用了递归的,从第1个结点开始的。
printf("%d\n", ans);
return 0;
}