[Hackerearth] EvenOddity

 

Problem Link    : EvenOddity
Category        : Data Structure, Counting, Bits, Centroid Decomposition
Contest         : December Circuits'23  
Tutorial        : Counting Subarrays 

#include "bits/stdc++.h"
using namespace std;
#define int     long long int
#define endl    '\n'
 
const int maxn = 1e6 + 5;
 
struct Node {
    int v, w;
    Node(int v = 0, int w = 0)
        : v(v), w(w) {}
};
 
int n;
vector<vector<Node>> graph;
vector<int> subtree;
vector<bool> isCentroid;
int curTreeNodes;
 
void calcSubtreeSize(int u, int p) {
    subtree[u] = 1;
    for (Node it : graph[u]) {
        int v = it.v;
        if (v == p or isCentroid[v]) {
            continue;
        }
        calcSubtreeSize(v, u);
        subtree[u] += subtree[v];
    }
}
 
int getCentroid(int u, int p) {
    for (Node it : graph[u]) {
        int v = it.v;
        if (v == p or isCentroid[v]) {
            continue;
        }
        if (subtree[v] > curTreeNodes / 2) {
            return getCentroid(v, u);
        }
    }
    return u;
}
 
int cache[2][2]; // cache[ distance % 2 ][ parity(xorSum) % 2 ]
int curCache[2][2];
int ans;
 
void add(int u, int p, int lvl, int xorSum) {
    int bit = lvl & 1;
    int parity = __builtin_popcount(xorSum) & 1;
    int curAns = cache[bit ^ 1][parity];
    ans += curAns;
    curCache[bit][parity]++;
    for (Node it : graph[u]) {
        int v = it.v;
        int w = it.w;
        if (v == p or isCentroid[v]) {
            continue;
        }
        add(v, u, lvl + 1, xorSum ^ w);
    }
}
 
void decompose(int u) {
    calcSubtreeSize(u, 0);
    curTreeNodes = subtree[u];
    int centroid = getCentroid(u, 0);
    isCentroid[centroid] = true;
    memset(cache, 0, sizeof cache);
    cache[0][0] = 1;
    for (Node it : graph[centroid]) {
        int child = it.v;
        if (isCentroid[child]) {
            continue;
        }
        memset(curCache, 0, sizeof curCache);
        add(child, centroid, 1, it.w);
        for (int i = 0; i < 2; i++) {
            for (int j = 0; j < 2; j++) {
                cache[i][j] += curCache[i][j];
            }
        }
    }
    for (Node it : graph[centroid]) {
        int child = it.v;
        if (isCentroid[child]) {
            continue;
        }
        decompose(child);
    }
}
 
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cout.precision(12);
 
    bool FILEIO = 1;
    if (FILEIO and fopen("f2.txt", "r")) {
        freopen("f2.txt", "r", stdin);
        freopen("f2out.txt", "w", stdout);
    }
 
    int tc;
    cin >> tc;
    for (int tcase = 1; tcase <= tc; tcase++) {
        cin >> n;
        graph.clear(); graph.resize(n + 1);
        subtree.clear(); subtree.resize(n + 1);
        isCentroid.clear(); isCentroid.resize(n + 1);
        for (int e = 1; e < n; e++) {
            int u, v, w;
            cin >> u >> v >> w;
            graph[u].emplace_back(v, w);
            graph[v].emplace_back(u, w);
        }
        ans = 0;
        decompose(1);
        cout << ans * 2 << endl;
    }
}

[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;  
}