Scribd
Upload
71 views2 pages

Dsu On Tree

This document contains C++ code for counting the number of nodes of each color in the subtrees of a tree. It uses recursion to traverse the tree depth-first and maintain arrays cnt and v to track the counts and nodes in each subtree. The dfs function is called recursively on each node, first processing its "big child" subtree if it exists, then the remaining "small child" subtrees, updating cnt and v appropriately at each level.

Uploaded by

prajjwal
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
71 views2 pages

Dsu On Tree

This document contains C++ code for counting the number of nodes of each color in the subtrees of a tree. It uses recursion to traverse the tree depth-first and maintain arrays cnt and v to track the counts and nodes in each subtree. The dfs function is called recursively on each node, first processing its "big child" subtree if it exists, then the remaining "small child" subtrees, updating cnt and v appropriately at each level.

Uploaded by

prajjwal
Copyright
© © All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
You are on page 1/ 2

const int N=500005;

vector<int>g[N];
vector<int>v[N];// store all the node which are in the subtree of s
int sz[N];
int cnt[N];//cnt[i] is count of node in subtree whose color is i
int color[N];
void getsz(int s, int p)
{
sz[s]=1;
for(auto u:g[s])
{
if(u==p)
continue;
getsz(u, s);
sz[s]+=sz[u];
}
}
//keep==1 we are in bigChild
//keep==0 we are in smallChild
void dfs(int s, int p, bool keep) {

int Max = -1, bigchild = -1;


for (auto u : g[s]) {
if (u != p && Max < sz[u]) {
Max = sz[u];
bigchild = u;
}
}

for (auto u : g[s]) {


if (u != p && u != bigchild) {
dfs(u, s, 0);
}
}

if (bigchild != -1) {
dfs(bigchild, s, 1);
// at this moment the cnt[N] array store the data of BigChild
// which is reuse in the calculation of subtree s(same for v)
//so we just swap the data.
swap(v[s], v[bigchild]);
}

v[s].push_back(s);
cnt[color[s]]++;

//to get data from small subtree we just iterate through node
for (auto u : g[s]) {
if (u != p && u != bigchild) {
for (auto x : v[u]) {
cnt[color[x]]++;
v[s].push_back(x);
}
}
}
// now for the subtree s cnt[i] array store the number of node with color i
//v[s] store all the elemnet in the subtree of s.

// if the keep=1 then this data is being use in the calculation of the parent
ans
//But for keep=0 we have to erase data from cnt array.
if (keep == 0) {
for (auto u : v[s]) {
cnt[color[u]]--;
}
}
}

You might also like