[Codechef] Prime Distance on Tree

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : Prime Distance on Tree
Source            : Codechef
Category          : Graph Theory
Algorithm         : Centroid Decomposition
Verdict           : Accepted

#include "bits/stdc++.h"  
  
using namespace std;  
  
#define ll          long long int  
  
static const int maxn = 1e5 + 5;  
static const int logn = 18;  
  
vector <int> graph[maxn];  
int subtreeSize[maxn];  
bool is_centroid[maxn];  
int num;  
  
ll ndist[logn][maxn];  
  
bool isPrime[maxn];  
vector <int> prime;  
  
inline void seive()  
{  
      for (int i = 0; i < maxn; i++) isPrime[i] = 1;  
      isPrime[0] = isPrime[1] = 0;  
      for (int i = 4; i < maxn; i += 2) isPrime[i] = 0;  
      for (int i = 3; i * i <= maxn; i += 2)  
      {  
            if (isPrime[i])  
            {  
                  for (int j = i*i; j < maxn; j += i+i) isPrime[j] = 0;  
            }  
      }  
      prime.emplace_back(2);  
      for (int i = 3; i < maxn; i += 2) if (isPrime[i]) prime.emplace_back(i);  
}  
  
inline void dfs(int u, int p = -1)  
{  
      subtreeSize[u] = 1;  
      for (int v : graph[u])  
      {  
            if (v == p or is_centroid[v]) continue;  
            dfs(v, u);  
            subtreeSize[u] += subtreeSize[v];  
      }  
}  
  
inline int get_centroid(int u, int p = -1)  
{  
      for (int v : graph[u])  
      {  
            if (v == p or is_centroid[v]) continue;  
            if (subtreeSize[v] > (num/2)) return get_centroid(v, u);  
      }  
      return u;  
}  
  
inline void Add(int u, int p, int depth, int dist, int add)  
{  
      ndist[depth][dist] += add;  
      for (int v : graph[u])  
      {  
            if (v == p or is_centroid[v]) continue;  
            Add(v, u, depth, dist+1, add);  
      }  
}  
  
inline ll calcPrimeDistance(int u, int p, int depth, int dist)  
{  
      ll ret = 0;  
      for (int p : prime)  
      {  
            if (p < dist) continue;  
            if (ndist[depth][p-dist] == 0) break;  
            if (p != dist) ret += ndist[depth][p-dist];  
            else ret += 2*ndist[depth][p-dist];  
      }  
      for (int v : graph[u])  
      {  
            if (v == p or is_centroid[v]) continue;  
            ret += calcPrimeDistance(v, u, depth, dist+1);  
      }  
      return ret;  
}  
  
ll ans;  
  
inline void decompose(int u, int depth)  
{  
      dfs(u, -1);  
      num = subtreeSize[u];  
      int centroid = get_centroid(u);  
      is_centroid[centroid] = 1;  
      Add(centroid, -1, depth, 0, 1);  
      ll ret = 0;  
      for (int v : graph[centroid])  
      {  
            if (is_centroid[v]) continue;  
            Add(v, centroid, depth, 1, -1);  
            ret += calcPrimeDistance(v, centroid, depth, 1);  
            Add(v, centroid, depth, 1, 1);  
      }  
      ans += ret / 2;  
      for (int v : graph[centroid])  
      {  
            if (is_centroid[v]) continue;  
            decompose(v, depth+1);  
      }  
      for (int i = 0; i < maxn && ndist[depth][i]; i++) ndist[depth][i] = 0;  
}  
  
int main()  
{  
      seive();  
      int tnode;  
      cin >> tnode;  
      for (int e = 1; e < tnode; e++)  
      {  
            int u, v;  
            cin >> u >> v;  
            graph[u].emplace_back(v);  
            graph[v].emplace_back(u);  
      }  
      decompose(1, 0);  
      ll nCr = (ll)tnode * (ll)(tnode-1) / 2;  
      double pro = (double)ans / (double)nCr;  
      cout << fixed << setprecision(6) << pro << endl;  
}