[Hackerearth] Prefix queries

Algorithm: Persistent Trie

#include "bits/stdc++.h"
using namespace std;
#define int     long long int
#define endl    '\n'
 
const int maxn = 3e5 + 5;
const int maxbit = 62;
const int MAXNODES = 2e7;
 
// https://www.hackerearth.com/practice/data-structures/advanced-data-structures/trie-keyword-tree/practice-problems/algorithm/prefix-queries-3-2c111491/
struct Node {
    int st, lchild, rchild;
    Node(int st = 0, int lchild = 0, int rchild = 0)
        : st(st), lchild(lchild), rchild(rchild) {}
};
 
int NODES;
Node Tree[MAXNODES];
int version[maxn];
 
int update(int root, int pos, int val, int add) {
    int newNode = ++NODES;
    Tree[newNode] = Tree[root];
    if (pos < 0) {
        Tree[newNode].st += add;
        return newNode;
    }
    bool bit = val >> pos & 1LL;
    if (bit == 0) {
        Tree[newNode].lchild = update(Tree[newNode].lchild, pos - 1, val, add);
    } else {
        Tree[newNode].rchild = update(Tree[newNode].rchild, pos - 1, val, add);
    }
    Tree[newNode].st = Tree[ Tree[newNode].lchild ].st + Tree[ Tree[newNode].rchild ].st;
    return newNode;
}
 
int getMaxPrefix(int root, int pos, int val) {
    if (pos < 0) {
        return 0;
    }
    bool bit = val >> pos & 1LL;
    if (bit == 0 and Tree[ Tree[root].lchild ].st > 0) {
        return 1 + getMaxPrefix(Tree[root].lchild, pos - 1, val);
    } else if (bit == 1 and Tree[ Tree[root].rchild ].st > 0) {
        return 1 + getMaxPrefix(Tree[root].rchild, pos - 1, val);
    }
    return 0;
}
 
int n;
int q;
vector<int> arr;
vector<vector<int>> graph;
vector<vector<pair<int, int>>> queries;
vector<vector<int>> level;
vector<int> par;
vector<int> ans;
 
void dfs(int u = 1, int p = 0, int lvl = 0) {
    version[1] = update(version[1], maxbit - 1, arr[u], 1);
    for (auto it : queries[u]) {
        ans[it.second] = getMaxPrefix(version[1], maxbit - 1, it.first);
    }
    for (int v : graph[u]) {
        if (v == p) {
            continue;
        }
        dfs(v, u);
    }
    version[1] = update(version[1], maxbit - 1, arr[u], -1);
}
 
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cout.precision(12);
 
    bool FILEIO = 1;
    if (FILEIO and fopen("in.txt", "r")) {
        freopen("in.txt", "r", stdin);
    }
 
    int tc;
    cin >> tc;
    while (tc--) {
        cin >> n;
        arr.clear(); arr.resize(n + 1);
        NODES = 0;
        for (int i = 1; i <= n; i++) {
            cin >> arr[i];
        }
        graph.clear(); graph.resize(n + 1);
        for (int e = 1; e < n; e++) {
            int u, v;
            cin >> u >> v;
            graph[u].push_back(v);
            graph[v].push_back(u);
        }
        cin >> q;
        queries.clear(); queries.resize(n + 1);
        for (int i = 0; i < q; i++) {
            int v, x;
            cin >> v >> x;
            queries[x].push_back({v, i});
        }
        ans.resize(q);
        dfs();
        for (int x : ans) {
            cout << x << ' ';
        }
        cout << endl;
 
        for (int i = 0; i <= NODES; i++) {
            Tree[i] = Node();
        }
    }
}