Merge pull request matthewsamuel95#872 from amulyagaur/patch-1

A problem on sack (https://www.hackerrank.com/challenges/coloring-tre…

Author: matthewsamuel95

Graph/Sack/sack(coloring-tree-hackerrank).cpp

@@ -0,0 +1,71 @@
+#include <bits/stdc++.h>
+using namespace std;
+set<int> *s[100005];
+vector< vector<int> > g(100005);
+int color[100005];
+int size_tree[100005];
+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)
+    {
+        s[u]=s[mv];   
+    }
+    else
+    {
+        s[u] = new set<int>();
+    }
+    (*s[u]).insert(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=(*s[v]).begin();it!= (*s[v]).end();it++)
+            {
+                (*s[u]).insert(*it);
+            }
+        }
+    }
+    return size_tree[u]=(int)(*s[u]).size();
+}
+
+int main()
+{
+    int n,m,r;
+    scanf("%d %d %d",&n,&m,&r);
+    
+    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);
+    }
+    for(int i=1;i<=n;i++) scanf("%d",&color[i]);
+    int sz=dfs(r,0);
+    while(m--)
+    {
+        int d;
+        scanf("%d",&d);
+        printf("%d\n",size_tree[d]);
+    }
+    return 0;  
+}