1923E - Count Paths - CodeForces Solution

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <sstream>
#include <queue>
#include <deque>
#include <bitset>
#include <iterator>
#include <list>
#include <stack>
#include <map>
#include <set>
#include <functional>
#include <numeric>
#include <utility>
#include <limits>
#include <time.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>
#include<cctype>
#include<fstream>
#include<set>
#include<sstream>
#include<memory>

using namespace std;

#define mem(t, v)           memset ((t) , v, sizeof(t))
#define all(x)              x.begin(),x.end()
#define un(x)               x.erase(unique(all(x)), x.end())
#define sf(n)               scanf("%d", &n)
#define sff(a,b)            scanf("%d %d", &a, &b)
#define sfff(a,b,c)         scanf("%d %d %d", &a, &b, &c)
#define sl(n)               scanf("%lld", &n)
#define sll(a,b)            scanf("%lld %lld", &a, &b)
#define slll(a,b,c)         scanf("%lld %lld %lld", &a, &b, &c)
#define D(x)                cerr << __LINE__ << ": " << #x " = " << (x) << '\n'
#define DD(x, y)            cerr << __LINE__ << ": " << #x " = " << (x) << ", " << #y " = " << (y) << '\n' 
#define DDD(x, y, z)        cerr << __LINE__ << ": " << #x " = " << (x) << ", " << #y " = " << (y) << ", " << #z " = " << (z) << '\n'

#define DBG         cerr << __LINE__ << ": " <<  "Hi" << '\n'
#define PI          acos(-1.00)
#define xx          first
#define yy          second
#define eps         1e-9

typedef long long int LL;
typedef double db;
typedef pair<int,int> pii;
typedef pair<LL,LL> pll;

#define MAX   200000

int color[MAX+1];

vector<int>tree[MAX+1];
shared_ptr<map<int, int>>dp[MAX+1];

LL ans;

int dfs(int u, int p)
{
  int bigchild = -1, big_size = 0;
  int my_size = 1;
  for(auto v: tree[u])
  {
    if(v!=p)
    {
      int cur_size = dfs(v, u);
      if(cur_size > big_size)
        bigchild = v, big_size = cur_size;
      my_size += cur_size;
    }
  }
  if(bigchild == -1)
  {
    dp[u] = make_shared<map<int,int>>();
    (*dp[u])[color[u]] = 1;
    return my_size;
  }
  else
  {
    dp[u] = dp[bigchild];
    ans += (*dp[u])[color[u]];
    (*dp[u])[color[u]] = 1;
    for(auto v: tree[u])
    {
      if(v != p and v != bigchild)
      {
        for(auto kvp: *dp[v])
        {
          int c = kvp.first;
          LL num = kvp.second;
          ans += (*dp[u])[c]*num;
          if(c != color[u])
            (*dp[u])[c] += num;
        }
      }
    }
    return my_size;
  }
}

LL solve()
{
  ans = 0;
  dfs(1, -1);
  return ans;
}


int main()
{
    //freopen("in.txt","r",stdin);
    //freopen("out.txt","w",stdout);
    int t;
    sf(t);
    for(int cs=1; cs<=t; cs++)
    {
      int n;
      sf(n);
      for(int i=0; i<n; i++)
        sf(color[i]);
      for(int i = 0; i<n-1; i++)
      {
        int u, v;
        sff(u, v);
        u--, v--;
        tree[u].push_back(v);
        tree[v].push_back(u);
      }
      //DBG;
      printf("%lld\n",solve());
      for(int i = 0; i<n; i++)
      {

        //D(dp[i]==nullptr);
        tree[i].clear();
        //dp[i].reset();
      }
    }
    return 0;
}