AcWing 4998. 暴搜-树的存储为y总模板-70分
原题链接
困难
作者:
yge
,
2024-04-06 17:05:30
,
所有人可见
,
阅读 48
C++ 代码
#include<bits/stdc++.h>
#include<map>
using namespace std;
const int N = 2e5 +10;
const int M = 2 * N;
int color[N];
int e[M],h[N],ne[M],idx=0;
bool st[N];
int ans = 0;
void add(int a, int b)
{
e[idx] = b; ne[idx] = h[a] ; h[a] = idx ++;
}
bool check(map<int , int> mp) //判定树u是不是平衡树
{
int pre = 0;
for(auto &i : mp)
{
if(pre && i.second != pre) return false;
pre = i.second;
}
return true;
}
map<int , int> dfs(int u)
{
st[u] = true;
map<int , int> sum ;
if(h[u] == -1) //如果是叶子结点
{
sum[color[u]]++;
ans ++ ;
return sum;
}
sum[color[u]]++;
for( int i = h[u]; i != -1; i = ne[i] )
{
int j = e[i]; //j是子树
if(st[j] == false)
{
map<int ,int > res = dfs(j);
for(auto&x : res)
{
sum[x.first] += x.second; //把子树颜色更新到u中
}
}
}
if(check(sum)) ans++;//检查自己是不是平衡树,如果是则答案加1
return sum;
}
int main()
{
int n;
cin >> n;
memset(h,-1,sizeof h);
for(int i = 1; i <= n; i ++)
{
int c,f;
cin >> c >> f;
if(f==0) add(f,i);
else add(i,f),add(f,i);
color[i] = c;
}
dfs(1);
cout << ans;
}