Tutorial: Tree Rerooting approach for Path XOR Sums

Hello everyone! I recently solved this problem and wanted to share a Tree Rerooting (In-Out DP) solution based on the contribution technique. Problem Link: 766E - Mahmoud and a xor trip

If you're looking for a structured way to handle path XOR queries across a tree, breaking it down into bits and using Tree DP is the way.

1. The Contribution Technique

Instead of computing the XOR sum for every single path $$$d(x, y)$$$, we can process the answer bit by bit. Since, $$$a[i] \lt = 1e6$$$, 20 bits are enough.

For a specific bit $$$b \in [0, 20]$$$, we can treat the node values $$$a_i$$$ as a binary sequence (either $$$0$$$ or $$$1$$$ depending on the $$$b$$$-th bit). Our goal for each bit $$$b$$$ reduces to: Find the number of paths in the tree where the sum of the $$$b$$$-th bit is odd (i.e., XOR sum is $$$1$$$).

If we can find this count, let's call it $$$res(b)$$$, then our final answer is simply:

2. Downwards Tree DP (In-DP)

Let's root the tree arbitrarily (e.g., at node $$$1$$$). For a given bit $$$b$$$, let $$$a(u)$$$ be the value of the $$$b$$$-th bit of node $$$u$$$. We define our states: $$$c(u, 0) =$$$ number of paths starting at $$$u$$$ and going down into its subtree with an XOR sum of $$$0$$$. $$$c(u, 1) =$$$ number of paths starting at $$$u$$$ and going down into its subtree with an XOR sum of $$$1$$$.

Base Case: A single node forms a valid path by itself.

Transitions: For every child $$$v$$$ of $$$u$$$, if a path from $$$v$$$ has a sum of $$$S$$$, extending it to $$$u$$$ makes the sum $$$S \oplus a(u)$$$. To get a total sum of $$$0$$$, we need $$$S \oplus a(u) = 0 \implies S = a(u)$$$. To get a total sum of $$$1$$$, we need $$$S \oplus a(u) = 1 \implies S = 1 \oplus a(u)$$$.

Therefore, we sum over all children $$$v$$$:

3. Upwards Tree Rerooting (Out-DP)

To get the paths that go up from $$$u$$$ (towards its parent and into other branches), we define: $$$up(u, 0) =$$$ number of paths starting at $$$u$$$ and going up with an XOR sum of $$$0$$$. $$$up(u, 1) =$$$ number of paths starting at $$$u$$$ and going up with an XOR sum of $$$1$$$.

When transitioning from $$$u$$$ to its child $$$v$$$, the paths going "up" from $$$v$$$ consist of all valid paths starting at $$$u$$$ except the ones that go down into $$$v$$$'s subtree.

Let $$$x = a(v)$$$ and $$$y = a(u)$$$. The rerooting transitions are:

Explanation: $$$up(u, x) + c(u, x)$$$ gives all paths starting at $$$u$$$ with XOR sum $$$x$$$. We subtract $$$c(v, x \oplus y)$$$ to remove the paths that strictly lie in $$$v$$$'s subtree. Adding $$$v$$$ to these paths toggles their sum by $$$x$$$, giving us the required $$$0$$$ or $$$1$$$ sum for $$$up(v, \dots)$$$.

4. Assembling the Answer for Bit $$$b$$$

Once both DP passes are complete, the total number of paths starting at any node $$$u$$$ with an XOR sum of $$$1$$$ is $$$c(u, 1) + up(u, 1)$$$.

However, if we simply sum this over all nodes:

We will double-count every path of length $$$ \gt 1$$$ (since it is counted once from each endpoint). Single-node paths (where $$$a(u) = 1$$$) are only counted once.

To fix this, we subtract the single-node paths, divide by 2, and then add the single-node paths back:

Finally, add $$$res(b) \cdot 2^b$$$ to your global answer!

Code for reference:

#include <bits/stdc++.h>
using namespace std;
 
#define fastio() ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL)
 
#define int long long
#define double long double
 
typedef long long ll;
typedef unsigned long long ull;
 
typedef vector<int> vi;
typedef vector<ll> vll;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
 
#define pb push_back
#define ff first
#define ss second
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
 
#define YES cout << "YES\n"
#define NO cout << "NO\n"
#define endl '\n'
 
#define IN_VEC(v,n) vector<int> v(n); for(auto &i : v) cin >> i;
#define IN_ARR(arr,n) for(int i=0;i<n;i++) cin >> arr[i];
 
const ll INF = 1e18;
const int MOD = 1e9 + 7;
void dfs(int u,int p,int b,vector<vi> &dp,vector<vi> &a,vector<vi> &adj){
	dp[a[b][u]][u]=1;
	for(int v:adj[u]){
		if(v==p) continue;
		dfs(v,u,b,dp,a,adj);
		dp[0][u] += dp[a[b][u]][v];
		dp[1][u] += dp[1-a[b][u]][v];
	}
}
void updp(int u,int p,int b,vector<vi> &dp,vector<vi> &up,vector<vi> &a,vector<vi> &adj){
	for(int v:adj[u]){
		if(v==p) continue;
		up[0][v] += up[a[b][v]][u] + dp[a[b][v]][u]-dp[a[b][v]^a[b][u]][v];
		int t=1-a[b][v];
		up[1][v] += up[t][u] + dp[t][u]-dp[t^a[b][u]][v];
		updp(v,u,b,dp,up,a,adj);
	}
}
void solve() {
    int n;
    cin>>n;
    vector<vi> a(20,vi(n,0));
    for(int i=0;i<n;++i){
    	int x;
    	cin>>x;
    	for(int j=0;j<20;++j){
    		if(x&(1ll<<j)) a[j][i]=1;
    	}
    }
    vector<vi> adj(n);
    for(int i=0;i<n-1;++i){
    	int u,v;
    	cin>>u>>v;
    	u--,v--;
    	adj[u].pb(v);
    	adj[v].pb(u);
    }
    int ans=0;
    vector<vi> dp,up;
    for(int b=0;b<20;++b){
    	dp.assign(2,vi(n,0ll));
    	up.assign(2,vi(n,0ll));
    	dfs(0,-1,b,dp,a,adj);
    	updp(0,-1,b,dp,up,a,adj);
    	int cnt=0;
    	for(int i=0;i<n;++i) cnt += dp[1][i] + up[1][i]-(a[b][i]==1);
    	cnt/=2;
    	for(int i=0;i<n;++i) if(a[b][i]==1) cnt++;
    	ans += (1ll<<b)*cnt;
    	dp.clear();
    	up.clear();
    }
    cout<<ans<<endl;
}
 
int32_t main() {
    fastio();
    solve();
    return 0;
}

Hope this makes the In-Out DP logic clear.