PROBLEM LINK:
Author: Sandeep Singh
Tester: Arkaprabho Ghosh
Editorialist: Sandeep Singh
DIFFICULTY:
Medium
PREREQUISITES:
dsu on trees, dfs
QUICK EXPLANATION:
Using dsu on trees we can keep track of the number of nodes belonging to a group in the subtree of a node in an array say count[] such that at any node count[k] gives the count of nodes belonging to group k in the subtree of that node. Thus the task is to keep track of max value in the count array for each node. We can do that by realizing the fact that for a node the max value is at least the max value for the biggest child of that node. After that, we check the count against the current max whenever there is an increment in the count array while merging the other children. The current max after all the children are merged is the answer for that node.
EXPLANATION:
This is a basic problem of dsu on trees. You can refer to this link to read about it. I suggest you to first read about dsu on trees and then follow the editorial else nothing will make much sense.
Note, bigchild of a node is the child of a node with max number of nodes in its subtree. Thus, we need to calculate size of subtree of each node first.
Assuming you have read about dsu on trees, we know that we can easily keep track of count of nodes belonging to each group in the subtree of a node. Dsu on trees is all about keeping the values of the biggest child of a node and merging it with other children. We do the same here.
For leaf nodes, the answer is obviously 1, for nonleaf, we first initialize the answer to the answer of the biggest child since subtree of the biggest child is also a part of the subtree of the node, hence the max count is at least the answer for bigchild. After this, while we merge the nodes of subtrees of all other children, at each step whenever we increment a count for a group, we compare it against the current maximum and make changes if needed. After all the nodes are merged, we store the answer in an array. In the end, we display that array.
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
#define ll long long int
#define max(a,b) (a>b?a:b)
#define min(a,b) (a<b?a:b)
#define scanarr(a,b,c) for(int ii=b;ii<c;ii++)cin>>a[ii]
#define showarr(a,b,c) for(int ii=b;ii<c;ii++)cout<<a[ii]<<' '
#define ln cout<<'\n'
#define FAST ios_base::sync_with_stdio(false);cin.tie(NULL);
#define mod 1000000007
#define MAX 50005
using namespace std;
////////////////////////////////////////////////////////////////CODE STARTS HERE////////////////////////////////////////////////////////////////
int prop[MAX],ans[MAX],cnt[MAX],sz[MAX];
vector<int>g[MAX];
vector<int> *vec[MAX];
void getsize(int u,int p){
sz[u]=1;
for(int child:g[u]){
if(child!=p){
getsize(child,u);
sz[u]+=sz[child];
}
}
}
void dfs(int u,int p,bool keep){
int mx=-1,bigChild=-1;
for(int child:g[u])
if(child!=p && sz[child]>mx)
bigChild=child,mx=sz[bigChild];
for(int child:g[u])
if(child!=p && child!=bigChild)
dfs(child,u,0);
if(bigChild!=-1)
dfs(bigChild,u,1),vec[u]=vec[bigChild],ans[u]=ans[bigChild];
else
vec[u]=new vector<int>();
vec[u]->push_back(u);
cnt[prop[u]]++;
ans[u]=max(cnt[prop[u]],ans[u]);
for(int child:g[u]){
if(child!=p && child!=bigChild){
for(int x:*vec[child]){
cnt[prop[x]]++;
ans[u]=max(cnt[prop[x]],ans[u]);
vec[u]->push_back(x);
}
}
}
if(keep==0){
for(int child:*vec[u]){
cnt[prop[child]]--;
}
}
}
void solve(){
int n,x,y;
cin>>n;
scanarr(prop,1,n+1);
for(int i=1;i<n;i++){
cin>>x>>y;
g[x].push_back(y);
g[y].push_back(x);
}
getsize(1,0);
dfs(1,0,1);
showarr(ans,1,n+1);
ln;
}
int main(){
FAST
solve();
}