This article is about a solution to a variation of [problem:894D], with (almost) complete binary trees replaced by arbitrary binary trees. The solution turned out to be more complicated than I thought, so I decided to post it as a separate blog. ↵
The main idea is to fix the LCA, just like the original problem. ↵
Note: We are calculating the sum of all $H_i - L$, such that $$H_i - L \geq 0 \iff H_i - (d_{A_i} + d_x - 2d_{lca(A_i,x)}) \geq 0 \iff H_i - d_{A_i} + 2d_{lca(A_i,x)} \geq d_x \iff d_x \leq H_i - d_{A_i} + 2d_{lca(A_i,x)}$$ ↵
where $x$ is any destination vertex and $d_v$ is the distance from the root to the $v$-th vertex. ↵
Therefore, we could calculate the answer using the following functions: ↵
↵
~~~~~↵
// dist[v] = sorted list of d_x, where x is any node in v's subtree↵
// prefsum[v] = prefix sum array of dist[v]↵
typedef long long i64;↵
i64 helper(int v, i64 val) {↵
if(v == 0) return 0;↵
auto x = std::upper_bound(dist[v].begin(), dist[v].end(), val);↵
if(x == dist[v].begin()) return 0;↵
i64 cnt = std::distance(dist[v].begin(), x);↵
return cnt * val - prefsum[v][cnt - 1];↵
}↵
i64 solve(int v, int c, i64 happiness, int start) { // par[root] = 0↵
if(v == 0) return 0;↵
i64 val = happiness - d[start] + 2 * d[v];↵
return helper(v, val) - helper(c, val) + solve(par[v], v, happiness, start);↵
}↵
i64 query(int v, i64 happiness) {↵
return solve(v, 0, happiness, v);↵
}↵
↵
~~~~~ ↵
We could use this code to calculate the answer for an arbitrary tree, however queries would take $O(depth \times \log{n}) = O(n \log{n})$ time. ↵
We can speed this algorithm up by using centroid decomposition. We will create a "centroid tree", and create a vector for each node that stores the distances from each node to the nodes in its subtree. ↵
Let $d(x,y)$ be the distance from node $x$ to node $y$. Then, we need to calculate the sum of all $H_i - L$ such that $H_i - L \geq 0 \iff H_i - (d(v,A_i) + d(v, x)) \geq 0 \iff d(v,x) \leq H_i - d(v,A_i)$, where $v$ is an ancestor of $A_i$ in the centroid tree, and $x$ is a node in $v$'s subtree, but *not* in the subtree containing $A_i$. ↵
Note that we can use the exact same helper function to calculate the answer for a single vertex, however we also need to find a way to subtract the answer for $v$'s children. We can reduce subtree queries to range queries on an array using the Euler tour technique. ↵
We also need a data structure that supports the following operations: ↵
1. $count(l, r, x)$: count the number of elements in [l, r] such that $a_i \leq x$. ↵
2. $sum(l, r, x)$: calculate the sum of elements in [l, r] such that $a_i \leq x$. ↵
We can use a wavelet tree/persistent segment tree/simple segment tree (offline). ↵
In the following code, let `idx[v]` be the index of vertex $v$ in the DFS order, and `sz[v]` be the size of $v$'s subtree in the centroid tree. Then, the code can be adapted for the centroid tree like this: ↵
↵
~~~~~↵
// we build a separate data structure for every vertex↵
data_structure DS[MAXN];↵
i64 helper(int v, int l, int r, int val) {↵
if(l > r) return 0;↵
return DS[v].count(l, r, val) * val - DS[v].sum(l, r, val);↵
}↵
i64 solve(int v, int c, int happiness, int start) {↵
if(v == 0) return 0;↵
int x = idx[c] - idx[v] + 1;↵
int y = x + sz[c] - 1;↵
i64 val = happiness - d[v][x];↵
return (c == 0 ? helper(v, 1, sz[v], val) : helper(v, 1, x - 1, val) + helper(v, y + 1, sz[v], val)) + solve(par[v], v, happiness, start);↵
}↵
i64 query(int v, i64 happiness) {↵
return solve(v, 0, happiness, v);↵
}↵
~~~~~↵
↵
This solution takes $O(\log^2{n})$ time per query. (`solve` is called $O(\log{n})$ times per query, and every `helper` query takes $O(\log{n})$ time.) ↵
*Disclaimer*: This is the first time I have used centroid decomposition. Feel free to point out any flaws in this approach. I will try to post a full solution later.↵
*Update*: The initial version stated that this was a solution for arbitrary binary trees. It is actually a solution for any tree.
The main idea is to fix the LCA, just like the original problem. ↵
Note: We are calculating the sum of all $H_i - L$, such that $$H_i - L \geq 0 \iff H_i - (d_{A_i} + d_x - 2d_{lca(A_i,x)}) \geq 0 \iff H_i - d_{A_i} + 2d_{lca(A_i,x)} \geq d_x \iff d_x \leq H_i - d_{A_i} + 2d_{lca(A_i,x)}$$ ↵
where $x$ is any destination vertex and $d_v$ is the distance from the root to the $v$-th vertex. ↵
Therefore, we could calculate the answer using the following functions: ↵
↵
~~~~~↵
// dist[v] = sorted list of d_x, where x is any node in v's subtree↵
// prefsum[v] = prefix sum array of dist[v]↵
typedef long long i64;↵
i64 helper(int v, i64 val) {↵
if(v == 0) return 0;↵
auto x = std::upper_bound(dist[v].begin(), dist[v].end(), val);↵
if(x == dist[v].begin()) return 0;↵
i64 cnt = std::distance(dist[v].begin(), x);↵
return cnt * val - prefsum[v][cnt - 1];↵
}↵
i64 solve(int v, int c, i64 happiness, int start) { // par[root] = 0↵
if(v == 0) return 0;↵
i64 val = happiness - d[start] + 2 * d[v];↵
return helper(v, val) - helper(c, val) + solve(par[v], v, happiness, start);↵
}↵
i64 query(int v, i64 happiness) {↵
return solve(v, 0, happiness, v);↵
}↵
↵
~~~~~ ↵
We could use this code to calculate the answer for an arbitrary tree, however queries would take $O(depth \times \log{n}) = O(n \log{n})$ time. ↵
We can speed this algorithm up by using centroid decomposition. We will create a "centroid tree", and create a vector for each node that stores the distances from each node to the nodes in its subtree. ↵
Let $d(x,y)$ be the distance from node $x$ to node $y$. Then, we need to calculate the sum of all $H_i - L$ such that $H_i - L \geq 0 \iff H_i - (d(v,A_i) + d(v, x)) \geq 0 \iff d(v,x) \leq H_i - d(v,A_i)$, where $v$ is an ancestor of $A_i$ in the centroid tree, and $x$ is a node in $v$'s subtree, but *not* in the subtree containing $A_i$. ↵
Note that we can use the exact same helper function to calculate the answer for a single vertex, however we also need to find a way to subtract the answer for $v$'s children. We can reduce subtree queries to range queries on an array using the Euler tour technique. ↵
We also need a data structure that supports the following operations: ↵
1. $count(l, r, x)$: count the number of elements in [l, r] such that $a_i \leq x$. ↵
2. $sum(l, r, x)$: calculate the sum of elements in [l, r] such that $a_i \leq x$. ↵
We can use a wavelet tree/persistent segment tree/simple segment tree (offline). ↵
In the following code, let `idx[v]` be the index of vertex $v$ in the DFS order, and `sz[v]` be the size of $v$'s subtree in the centroid tree. Then, the code can be adapted for the centroid tree like this: ↵
↵
~~~~~↵
// we build a separate data structure for every vertex↵
data_structure DS[MAXN];↵
i64 helper(int v, int l, int r, int val) {↵
if(l > r) return 0;↵
return DS[v].count(l, r, val) * val - DS[v].sum(l, r, val);↵
}↵
i64 solve(int v, int c, int happiness, int start) {↵
if(v == 0) return 0;↵
int x = idx[c] - idx[v] + 1;↵
int y = x + sz[c] - 1;↵
i64 val = happiness - d[v][x];↵
return (c == 0 ? helper(v, 1, sz[v], val) : helper(v, 1, x - 1, val) + helper(v, y + 1, sz[v], val)) + solve(par[v], v, happiness, start);↵
}↵
i64 query(int v, i64 happiness) {↵
return solve(v, 0, happiness, v);↵
}↵
~~~~~↵
↵
This solution takes $O(\log^2{n})$ time per query. (`solve` is called $O(\log{n})$ times per query, and every `helper` query takes $O(\log{n})$ time.) ↵
*Disclaimer*: This is the first time I have used centroid decomposition. Feel free to point out any flaws in this approach. I will try to post a full solution later.↵
*Update*: The initial version stated that this was a solution for arbitrary binary trees. It is actually a solution for any tree.
