Solving 894D for arbitrary trees in O(nlog^2n)
Difference between en2 and en3, changed 22 character(s)
This article is about a solution to a variation of [problem:894D], with (almost) complete binary trees replaced by arbitrary 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][
xidx[start] - idx[v] + 1];↵
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.

History

 
 
 
 
Revisions
 
 
  Rev. Lang. By When Δ Comment
en3 English f2lk6wf90d 2017-11-21 14:45:14 22 Tiny change: 'ss - d[v][x];\n retur' -> 'ss - d[v][idx[start] - idx[v] + 1];\n retur'
en2 English f2lk6wf90d 2017-11-21 00:16:10 138
en1 English f2lk6wf90d 2017-11-20 22:40:43 3906 Initial revision (published)