This article is about a solution to a variation of 894D - Ralph And His Tour in Binary Country, 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

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 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 , 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} ≤ *x*.

2. *sum*(*l*, *r*, *x*): calculate the sum of elements in [l, r] such that *a*_{i} ≤ *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][idx[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 time per query. (`solve`

is called times per query, and every `helper`

query takes 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.