HackerEarth Breaking Edges problem solution

In this HackerEarth Breaking Edges problem solution You are given a tree (undirected connected graph with no cycles) consisting of N nodes and N – 1 edges. There is a number associated with each node (vi) of the tree. You can break any edge of the tree you want and this would result in formation of 2 trees which were part of the original tree.

Let us denote by treeOr, the bitwise OR of all the numbers written on each node in a tree.

You need to find how many edges you can choose, such that, if that edge was removed from the tree, the treeOr of the 2 resulting trees is equal.

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 2e5 + 55;
vector<int> adj[MAXN];
int cnt[22][MAXN];
int POWER[22];
int dsu[MAXN + 55];
void DFS(int s, int par = -1) {
    for(auto it : adj[s]) {
        if(it == par) continue;
        DFS(it, s);
        for(int j = 0; j < 22; j++) {
            cnt[j][s] += cnt[j][it];
        }
    }
}
int ans = 0;
void dfs(int s, int par = -1) {
    for(auto it : adj[s]) {
        if(it == par) continue;
        dfs(it, s);
        bool can = true;
        for(int j = 0; j < 22 and can; j++) {
            if(cnt[j][1] != 0 and (cnt[j][1] == cnt[j][it] or cnt[j][it] == 0)) can = false;
        }
        if(can) ans++;
    }
}
int findRoot(int x) {
    if(dsu[x] == x) return x;
    return dsu[x] = findRoot(dsu[x]);
}
int main() {
    POWER[0] = 1;
    for(int i = 1; i < 22; i++) POWER[i] = POWER[i - 1] << 1;
    int n;
    cin >> n;
    assert(n >= 2 and n <= 200000);
    for(int i = 1; i <= n; i++) {
        dsu[i] = i;
        int val;
        scanf("%d", &val);
        for(int j = 0; j < 22; j++) if(POWER[j] & val) cnt[j][i]++;
        assert(val >= 0 and val <= 1048575);
    }
    for(int i = 1; i < n; i++) {
        int x, y;
        scanf("%d%d", &x, &y);
        adj[x].push_back(y);
        adj[y].push_back(x);
        int r1 = findRoot(x);
        int r2 = findRoot(y);
        dsu[r1] = r2;
        assert(r1 != r2);
        assert(x != y);
        assert(x >= 1 and x <= n and y >= 1 and y <= n);
    }
    set<int> components;
    for(int i = 1; i <= n; i++) components.insert(findRoot(i));
    assert((int)components.size() == 1);
    DFS(1);
    dfs(1);
    cout << ans << endl;
    return 0;
}