HackerEarth Dexters Random Generator problem solution

In this HackerEarth Dexter's Random Generator problem solution Dexter has recently constructed a random generator. This generator works on a tree of N nodes. The information about the tree has to be fed into the generator in order for it to start working.

Each node, u of the tree has a value A[u] associated with it. The generator works as follows:

It takes as input two integers, u and v. It then outputs two integers X and Y. X can be either A[u] or A[v], with equal probability. Now, the generator selects a node i randomly and with equal probability from the path u - v (including u and v) in the tree and outputs the value of Y as A[i].

Let the above process be denoted by: Process(u,v), which takes input a pair of integers and outputs a pair of integers (X, Y).

Now, Hannah has been invited to test the random generator constructed by Dexter. Hannah has Q questions. Each question is as follows:

Hannah selects two nodes u and v of the tree. She wants to know the maximum value of Query(u,v).

HackerEarth Dexter's Random Generator problem solution.

#include <bits/stdc++.h>
using namespace std;
const int N = 2e5 + 100;
const int LIM = 31;
const int LOGN = 20;
vector <int> g[N];

int a[N];
int trie[N * LIM] , to[N * LIM][2], ptr;
int root[N];

int insert(int num, int lev, int root) {
  int newRoot = ++ptr;
  int cur = newRoot;
  for(int i = LIM-1 ; i >= 0 ; i --) {
    int bit = (num >> i) & 1;
    to[cur][bit^1] = to[root][bit^1];
    to[cur][bit] = ++ptr;
    root = to[root][bit];
    cur = to[cur][bit];
    trie[cur] = lev;
  }
  return newRoot;
}

int query(int num, int lev, int root) {
  int cur = root;
  int ret = 0;
  for(int i = LIM-1 ; i >= 0 ; i --) {
    int bit = (num >> i) & 1;
    if(to[cur][bit ^ 1] && trie[to[cur][bit^1]] >= lev) {
      ret = ret * 2 + 1;
      cur = to[cur][bit ^ 1];
    }
    else {
      ret = ret * 2;
      cur = to[cur][bit];
    }
  }
  return ret;
}

int level[N] , P[LOGN][N];

void dfs(int u , int p){
  P[0][u] = p;
  for(int i = 1 ; i < LOGN ; i ++)
    P[i][u] = P[i-1][P[i-1][u]];
  root[u] = insert(a[u], level[u], root[p]);
  for(int i = 0 ; i < g[u].size();i++)
    if(g[u][i] != p){
      level[g[u][i]] = level[u] + 1;
      dfs(g[u][i] , u);
    }
}

inline int LCA(int x , int y){
  if(level[x] < level[y])swap(x , y);
  int lg = 1; for(;(1<<lg) <= level[x] ; lg ++);lg--;
  for(int i = lg ; i >= 0 ; i--)
    if(level[x] - (1<<i) >= level[y])
      x = P[i][x];
  if(x == y)return x;
  for(int i = lg ; i >= 0 ; i--)
    if(P[i][x] != -1 && P[i][x] != P[i][y]){
      x = P[i][x];
      y = P[i][y];
    }
  return P[0][x];
}

int main() {
  ptr = 1;
  int n,m; cin >> n >> m;
  for(int i = 1; i <= n; i ++) cin >> a[i];
  for(int i = 1; i < n; i ++) {
    int u,v; cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  dfs(1,0);
  while(m --) {
    int u,v; cin >> u >> v;
    int lca = LCA(u,v);
    int lev = level[lca]; 
    int ans = query(a[u], lev, root[u]);
    ans = max(ans, query(a[u], lev, root[v]));
    ans = max(ans, query(a[v], lev, root[u]));
    ans = max(ans, query(a[v], lev, root[v]));
    cout << ans << endl;
  }
  return 0;
}