Merge pull request matthewsamuel95#877 from amulyagaur/patch-5

Code for HEDWIG problem on CodeChef using sack technique

Author: matthewsamuel95

Graph/Sack/HEDWIG.cpp

@@ -0,0 +1,100 @@
+#include <bits/stdc++.h>
+using namespace std;
+map<int,int> *freq[200005];
+map<int,int> *counter[200005];
+vector< vector<int> > g(200005);
+int color[200005];
+int fav[200005];
+int ans[200005];
+int dfs(int u,int par)
+{
+    int mx=-1,mv=-1,v;
+    for(int i=0;i<g[u].size();i++)
+    {
+        v= g[u][i];
+        if(v==par)
+        continue;
+        int sz = dfs(v,u);
+        if(sz>mx)
+        {
+            mx=sz;
+            mv=v;
+        }
+    }
+    if(mx!=-1)
+    {
+        freq[u]=freq[mv];
+        counter[u]=counter[mv];   
+    }
+    else
+    {
+        freq[u] = new map<int,int>();
+        counter[u] = new map<int,int>();
+    }
+
+    if((*freq[u]).find(color[u]) != (*freq[u]).end())
+    {
+        (*counter[u])[(*freq[u])[color[u]]]--;    
+        if((*counter[u])[(*freq[u])[color[u]]]==0)
+        {
+           (*counter[u]).erase((*freq[u])[color[u]]);
+        }
+        (*freq[u])[color[u]]++;
+        (*counter[u])[(*freq[u])[color[u]]]++;
+    }
+    else
+    {
+        (*freq[u])[color[u]]++;
+        (*counter[u])[(*freq[u])[color[u]]]++;
+    }
+    for(int i=0;i<g[u].size();i++)
+    {
+        v= g[u][i];
+        if(v==par)
+        continue;
+        if(v!= mv && v!= par)
+        {
+            for(auto it=(*freq[v]).begin();it!= (*freq[v]).end();it++)
+            {
+                if((*freq[u]).find(it->first) != (*freq[u]).end())
+                {
+                    (*counter[u])[(*freq[u])[it->first]]--;
+                    if((*counter[u])[(*freq[u])[it->first]]==0)
+                    {
+                        (*counter[u]).erase((*freq[u])[it->first]);
+                    }
+                    (*freq[u])[it->first] += it->second;
+                    (*counter[u])[(*freq[u])[it->first]]++;
+                }
+                else
+                {
+                    (*freq[u])[it->first] = it->second;
+                    (*counter[u])[it->second]++;
+                }
+            }
+        }
+
+    }
+    if((*counter[u]).find(fav[u]) != (*counter[u]).end())
+    ans[u]= (*counter[u])[fav[u]];
+    return (int)(*freq[u]).size();
+}
+
+int main()
+{
+    int n;
+    scanf("%d",&n);
+    for(int i=1;i<=n;i++) scanf("%d",&color[i]);
+    for(int i=1;i<=n;i++) scanf("%d",&fav[i]);
+    for(int i=1;i<n;i++)
+    {
+        int u,v;
+        scanf("%d %d",&u,&v);
+        g[u].push_back(v);
+        g[v].push_back(u);
+    }
+    int sz=dfs(1,0);
+    for(int i=1;i<=n;i++)
+    printf("%d\n",ans[i]);
+    return 0;  
+}