[toph.co] Sofdor Ali's Black and White Tree

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : Sofdor Ali's Black and White Tree
Source            : toph.co
Category          : Graph Theory, Dynamic Programing
Algorithm         : DP on Tree
Verdict           : Accepted

#include <bits/stdc++.h>  
  
using namespace std;  
  
static const int maxn = 1e5 + 5;  
  
vector <int> graph[maxn];  
int color[maxn], white[maxn], black[maxn];  
int maxDiff = 0, maxWhite = -1, maxNode = -1;  
  
inline void dfs(int u = 0, int p = -1)  
{  
      if (color[u] == 0) black[u]++; // BLACK  
      else white[u]++;               // WHITE  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            dfs(v, u);  
            white[u] += white[v];  
            black[u] += black[v];  
      }  
      int diff = black[u] - white[u];  
      if (diff < 0)  
      {  
  
      }  
      else  
      {  
            if (diff > maxDiff)  
            {  
                  maxNode = u;  
                  maxDiff = diff;  
                  maxWhite = white[u];  
            }  
            else if (diff == maxDiff)  
            {  
                  if (maxWhite > white[u])  
                  {  
                        maxNode = u;  
                        maxWhite = white[u];  
                  }  
            }  
      }  
}  
  
int twhite = 0;  
  
inline void dfs2(int u = 0, int p = -1)  
{  
      if (u == maxNode)  
      {  
            twhite += black[u];  
            return;  
      }  
      else if (color[u] == 1)  
      {  
            twhite += 1;  
      }  
      for (int v : graph[u])  
      {  
            if (v == p) continue;  
            dfs2(v, u);  
      }  
}  
  
int main()  
{  
      //freopen("in.txt", "r", stdin);  
  
      ios_base::sync_with_stdio(false);  
      cin.tie(NULL);  
  
      int tnode;  
      cin >> tnode;  
      for (int i = 0; i < tnode; i++) cin >> color[i];  
      for (int i = 1; i < tnode; i++)  
      {  
            int u, v;  
            cin >> u >> v;  
            graph[u].push_back(v);  
            graph[v].push_back(u);  
      }  
      dfs();  
      dfs2();  
      cout << twhite << endl;  
}