题目描述
给定一棵树,结点由 $1$ 至 $n$ 编号,其中结点 $1$ 是树根。
树的每个点有一个颜色 $C_i$。
如果一棵树中存在的每种颜色的结点个数都相同,则我们称它是一棵颜色平衡树。
求出这棵树中有多少个子树是颜色平衡树。
输入格式
输入的第一行包含一个整数 $n$,表示树的结点数。
接下来 $n$ 行,每行包含两个整数 $C_i , F_i$,用一个空格分隔,表示第 $i$ 个结点的颜色和父亲结点编号。
特别地,输入数据保证 $F_1$ 为 $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
题目分析
两种做法
第一种就是树上莫队
分析题目,先暴力地想一下,其实这道题就是统计各个子树中符合各颜色数相同的个数
统计出现颜色的个数,这不就是莫队么
所以第一种做法直接对树进行前序遍历,预处理出前序遍历的树的颜色数组,再去对每个子树找到对应区间,这些区间就是询问,这样询问和原数组都有了,剩下的就是莫队模板了
这种情况下需要开$O(2)$优化
时间复杂度$O(n \sqrt{n})$ 需要卡常
#pragma GCC optimize(2)
#include <bits/stdc++.h>
using namespace std;
const int N = 200010 , M = N * 2;
struct Query
{
int l , r;
}q[N];
int n , len;
int h[N], e[M], ne[M], col[M], c[M], idx;
int s1[N] , s2[N] , l[N] , r[N]; // s1是颜色的个数,s2是当前颜色的个数相同的颜色数
void adde(int a, int b)
{
e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}
int get_len(int x)
{
return x / len;
}
int cidx = 0;
void dfs(int u)
{
col[++ cidx] = c[u];
l[u] = cidx;
for(int i = h[u] ; ~i ; i = ne[i])
dfs(e[i]);
r[u] = cidx;
}
void add(int x , int &cnt)
{
s2[s1[x]] --;
if(!s1[x]) cnt ++;
s1[x] ++;
s2[s1[x]] ++;
}
void del(int x , int &cnt)
{
s2[s1[x]] --;
s1[x] --;
if(!s1[x]) cnt --;
s2[s1[x]] ++;
}
int main()
{
scanf("%d", &n);
len = sqrt(n);
memset(h, -1, sizeof h);
int father;
for(int i = 1 ; i <= n ; i ++)
{
scanf("%d%d", &c[i], &father);
if(i != 1) adde(father , i);
}
dfs(1);
for (int i = 1 ; i <= n; i ++) {
q[i - 1] = {l[i] , r[i]};
}
sort(q , q + n , [&](const Query &a ,const Query &b){
if(get_len(a.l) != get_len(b.l))
return a.l < b.l;
return a.r < b.r;
});
int ans = 0, cnt = 0;
for (int k = 0, i = 0, j = 0; k < n; k ++) {
while(i < q[k].r) add(col[++ i] , cnt);
while(i > q[k].r) del(col[i --] , cnt);
while(j < q[k].l) del(col[j ++] , cnt);
while(j > q[k].l) add(col[-- j] , cnt);
int length = i - j + 1;
if(length % cnt == 0 && s2[length / cnt] == cnt) ans++; // 当前颜色的个数相同的颜色数等于当前颜色数
}
printf("%d" , ans);
return 0;
}
第二种是树上启发式合并,其实学过的可以发现,这道题和 AcWing 3189. Lomsat gelral 几乎一样,只是改变了询问的结果对象而已。
思路:树上启发式合并(DSU on tree)
1.先处理子树大小,类似树链刨分分一下重儿子和轻儿子,为了后续操作降低复杂度
2.遍历树,对于任何一棵树,先遍历轻儿子,如果有重儿子最后遍历,数据返回的时候只保留重儿子信息,更新颜色出现次数
3.如果子树颜色出现的最大数 * 最大数次数 = 子树的大小,那么这颗子树就是颜色平衡树,ans ++即可
板子几乎不变,改一下update()中的更新结果即可
时间复杂度$O(nlogn)$ 不需要卡常
#include <iostream>
#include <cstring>
using namespace std;
const int N = 200010 , M = N * 2;
int n;
int h[N], e[M], ne[M], idx;
int color[N], sz[N], son[N], cnt[N];
int res, sum;
int mx;
void add(int a, int b) {
e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}
int dfs_son(int u , int father) {
sz[u] = 1;
for (int i = h[u]; ~i; i = ne[i]) {
int j = e[i];
if(j == father) continue;
sz[u] += dfs_son(j , u);
if (sz[j] > sz[son[u]]) son[u] = j;
}
return sz[u];
}
void update(int u, int father , int val, int pson) {
int c = color[u];
cnt[c] += val;
if (cnt[c] > mx) mx = cnt[c], sum = 1;
else if (cnt[c] == mx) sum ++ ;
for (int i = h[u]; ~i; i = ne[i]) {
int j = e[i];
if (j == pson || j == father) continue;
update(j, u , val, pson);
}
}
void dfs(int u, int father , int op) {
for (int i = h[u]; ~i; i = ne[i]) {
int j = e[i];
if (j == son[u] || j == father) continue;
dfs(j, u , 0);
}
if (son[u]) dfs(son[u], u , 1);
update(u, father , 1, son[u]);
if (sum * mx == sz[u]) res ++ ;
if (!op) update(u, father , -1, 0), mx = 0;
}
int main() {
cin >> n;
memset(h, -1, sizeof h);
for (int i = 1; i <= n; i ++ ) {
int p;
cin >> color[i] >> p;
if(i != 1)add(p, i) , add(i , p);
}
dfs_son(1 , -1);
dfs(1, -1 , 1);
cout << res << endl;
return 0;
}