dsu on tree updated

Author: Nadim-Mahmud

Source/Trees/DSU on tree.cpp

@@ -1,70 +1,99 @@
-    ///1153 - Internet Bandwidth my code:
-vector <int> ed[102];
-int flw[102][102],cap[102][102],par[102],rcp[102];
-char vis[101];
-int s,t,d;
-int dfs()
-{
-    queue <int> q;
-    int v,f,i;
-    q.push(s);
-    MS(vis,0);
-    vis[s]=1;
-    rcp[s]=1E7;
-    while(!q.empty()){
-        f=q.front();
-        q.pop();
-        for(i=0;i<ed[f].size();i++){
-            v=ed[f][i];
-            if(!vis[v]&&(cap[f][v]-flw[f][v]>0)){
-                vis[v]=1;
-                par[v]=f;
-                rcp[v]=min(cap[f][v]-flw[f][v],rcp[f]);
-                if(v==d){
-                    return 1;
-                }
-                q.push(v);
-            }
+#include <bits/stdc++.h>
+#define ll long long
+#define P(X) cout<<"db "<<X<<endl;
+#define P2(X,Y) cout<<"db2 "<<X<<" "<<Y<<endl;
+#define rep(i,n) for(i=1;i<=n;i++)
+#define FO freopen("t.txt","w",stdout);
+#define MS(XX,YY) memset(XX,YY,sizeof(XX));
+#define pii pair<int,int>
+#define chk(n,i) (bool)(n&(1<<i))
+#define on(n,i) (n|(1<<i))
+#define off(n,i) n=n&(~(1<<i))
+#define eps 10e-7
+#define MX 1000005
+using namespace std;
+
+/**
+*  Description : Useful for all same types of operations on each node on a tree
+*               cnt[] storing all number of color of a type in sub tree of v
+*               res[] storing sum of color which apears maximum numbers of time
+*
+*  Time complexity : O(nlogn)
+*  Source :https://codeforces.com/blog/entry/44351
+*/
+
+
+vector<int>g[MX];
+
+int sz[MX];
+void getsz(int v, int p){
+    sz[v] = 1;  // every vertex has itself in its subtree
+    for(auto u : g[v])
+        if(u != p){
+            getsz(u, v);
+            sz[v] += sz[u]; // add size of child u to its parent(v)
         }
-    }
-    return 0;
 }
 
+ll cnt[MX],col[MX],ans[MX],res[MX],mxn;
+bool big[MX];
+
+void add(int v, int p, int x){
+    //ans[cnt[col[v]]] -= col[v];
+    cnt[col[v]] += x;
+    //ans[cnt[col[v]]] += col[v];
+    //mxn = max(cnt[col[v]],mxn);
+    for(auto u: g[v])
+        if(u != p && !big[u])
+            add(u, v, x);
+}
+
+
+void dfs(int v, int p, bool keep){
+    int mx = -1, bigChild = -1;
+    for(auto u : g[v])
+       if(u != p && sz[u] > mx)
+          mx = sz[u], bigChild = u;
+    //run a dfs on small childs and clear them from cnt
+    for(auto u : g[v])
+        if(u != p && u != bigChild)
+            dfs(u, v, 0);  //
+    // actual processing of vertex v starts from here
+    //mxn = 0;
+    if(bigChild != -1)
+        dfs(bigChild, v, 1), big[bigChild] = 1;  // bigChild marked as big and not cleared from cnt
+    // calculating ans
+    add(v, p, 1);
+    //res[v] = ans[mxn];
+    /** here access the result for each node. if needed then access on add() function
+        now cnt[c] is the number of vertices in subtree of vertex v that has color c.
+        You can answer the queries easily.
+    */
+    if(bigChild != -1)
+        big[bigChild] = 0;
+    if(keep == 0)
+        add(v, p, -1);
+}
+
+
+
 int main()
 {
-    int i,j,a,b,ts,cn=0,n,m,w,cp,cfl,tf;
+    ll i,j,a,b,ts,cn=0,cas=0,n,m,x,y,sum=0,mn=INT_MAX,u,v;
     //freopen("test.txt","r",stdin);
-    scanf("%d",&ts);
-    while(ts--){
-        scanf("%d %d %d %d",&n,&s,&d,&m);
-        for(i=0;i<=n;i++){
-            ed[i].clear();
-        }
-        MS(cap,0)
-        for(i=0;i<m;i++){
-            scanf("%d %d %d",&a,&b,&w);
-            if(w==0)continue;
-            if(!cap[a][b]){
-                ed[a].push_back(b);
-                ed[b].push_back(a);
-            }
-            cap[a][b]+=w;
-            cap[b][a]+=w;
-        }
-        tf=0;
-        MS(flw,0);
-        while(dfs()){
-            cp=d;
-            cfl=rcp[d];
-            tf+=cfl;
-            //P(tf)
-            do{
-                flw[par[cp]][cp]+=cfl;
-                flw[cp][par[cp]]-=cfl;
-                cp=par[cp];
-            }while(cp!=s);
-        }
-        printf("Case %d: %d\n",++cn,tf);
+    cin>>n;
+    for(int i = 1; i <= n; i++) {
+		cin >> col[i];
+	}
+    for(i=1;i<n;i++){
+        cin>>u>>v;
+        g[u].push_back(v);
+        g[v].push_back(u);
     }
+    getsz(1,0);
+    dfs(1,0,1);
+    for(int i = 1; i <= n; i++) {
+		cout << res[i] << ' ';
+	}
     return 0;
 }