PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Authors: deepak_changoi, iceknight1093
Testers: iceknight1093, mexomerf
Editorialist: iceknight1093
DIFFICULTY:
2699
PREREQUISITES:
Virtual/auxiliary trees, offline path add updates
PROBLEM:
Given a tree where each vertex has some color, define the value of an edge to be the sum of all the colors it disconnects upon removal.
Find the value of every edge.
EXPLANATION:
Let’s consider each color separately and see which edges it affects: if we are able to do this quickly enough for each edge, that will solve the problem.
For convenience, let’s root the tree at vertex 1.
Fix a color c, and let the vertices with color c be x_1, x_2, \ldots, x_k.
Note that the edges that disconnect c are exactly those edges that lie on the path between some x_i and some x_j, so we’d like to find all such edges quickly.
It’s easy to see that the set of these edges themselves forms a tree structure, which is called the auxiliary tree of this set of vertices.
Let T_c denote the auxiliary tree corresponding to color c.
It’s possible to build the auxiliary tree of \{x_1, x_2, \ldots, x_k\} using \mathcal{O}(k) LCA computations: the details of this may be found here or here.
This allows us to construct the auxiliary tree of every single color in \mathcal{O}(N\log N) or \mathcal{O}(N) time in total.
Now, consider some virtual tree T_c. We want to add c to every edge that lies in it.
While this is possible using tree decomposition techniques such as HLD combined with a lazy segment tree, the constraints of the problem are rather high and such solutions most likely won’t pass in time.
Instead, notice that all of our ‘updates’ are done before any ‘queries’ (where each query is the final value of an edge).
This allows us to generalize the trick of processing subarray addition updates offline.
What is this trick?
Suppose we have an array A, and we’d like to support Q updates of the form (L, R, x), each telling us to add x to the subarray [L, R].
To do this, we create a new array B, then add x to B_L and subtract x from B_{R+1}.
Finally, take the prefix sums of B: after this, B_i will denote the value that needs to be added to position i.
Instead of subarrays, we’ll use paths from a vertex to its ancestor.
Let’s create a new array B, with B_i = 0 for each 1 \leq i \leq N initially.
Now, consider some u \in T_c. Let p \in T_c be its parent (note that p is the parent of u in the auxiliary tree).
To add c to the path between u and p, we add c to B_u and subtract c from B_p.
Now, compute the subtree sums of every vertex using a dfs. Let S_u denote the subtree sum of vertex u.
It’s not hard to see that S_u is the answer for the edge connecting it to its parent!
Building the virtual trees takes \mathcal{O}(N\log N) time, and they give us \mathcal{O}(N) updates in total after which we perform a simple DFS.
So, we’ve solved the problem in \mathcal{O}(N\log N) time in total.
A few other problems that use virtual trees are:
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Setter's code (C++)
#include "bits/stdc++.h"
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
using namespace std;
using ll = long long int;
mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());
template<class T>
struct RMQ {
vector<vector<T>> jmp;
RMQ(const vector<T>& V) : jmp(1, V) {
for (int pw = 1, k = 1; pw * 2 <= (int)size(V); pw *= 2, ++k) {
jmp.emplace_back(size(V) - pw * 2 + 1);
for (int j = 0; j < (int)size(jmp[k]); ++j)
jmp[k][j] = min(jmp[k - 1][j], jmp[k - 1][j + pw]);
}
}
T query(int a, int b) {
assert(a < b); // or return inf if a == b
int dep = 31 - __builtin_clz(b - a);
return min(jmp[dep][a], jmp[dep][b - (1 << dep)]);
}
};
struct LCA {
int T = 0;
vector<int> time, out, dep, path, ret;
RMQ<int> rmq;
LCA(vector<vector<int>>& C) : time(size(C)), out(size(C)), dep(size(C)), rmq((dfs(C,0,-1), ret)) {}
void dfs(vector<vector<int>>& C, int v, int par) {
time[v] = T++;
for (int y : C[v]) if (y != par) {
path.push_back(v), ret.push_back(time[v]);
dep[y] = 1 + dep[v];
dfs(C, y, v);
}
out[v] = T;
}
int lca(int a, int b) {
if (a == b) return a;
tie(a, b) = minmax(time[a], time[b]);
return path[rmq.query(a, b)];
}
};
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector<int> col(n);
vector adj(n, vector<int>());
vector<array<int, 2>> edges;
map<int, vector<int>> vertices;
for (int i = 0; i < n; ++i) {
cin >> col[i];
vertices[col[i]].push_back(i);
}
for (int i = 0; i < n-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
edges.push_back({u, v});
}
LCA L(adj);
vector<ll> ans(n);
auto upd = [&] (int x, int y, int c) { // Add c to the (x, y) path, where x is an ancestor of y
ans[y] += c;
ans[x] -= c;
};
for (auto &[c, vlist] : vertices) {
// Build virtual tree of vertices with color c, adding to appropriate paths along the way
sort(begin(vlist), end(vlist), [&] (int u, int v) {return L.time[u] < L.time[v];});
int k = size(vlist);
for (int i = 0; i+1 < k; ++i) vlist.push_back(L.lca(vlist[i], vlist[i+1]));
sort(begin(vlist), end(vlist), [&] (int u, int v) {return L.time[u] < L.time[v];});
vlist.erase(unique(begin(vlist), end(vlist)), end(vlist));
stack<int> st;
for (int x : vlist) {
while (!st.empty()) {
int u = st.top();
if (L.out[u] >= L.out[x] and u != x) break;
st.pop();
}
if (!st.empty()) {
int u = st.top(); // u is the parent of x in this virtual tree
upd(u, x, c);
}
st.push(x);
}
}
auto dfs = [&] (const auto &self, int u, int p) -> void {
for (int v : adj[u]) {
if (v == p) continue;
self(self, v, u);
ans[u] += ans[v];
}
};
dfs(dfs, 0, 0);
for (auto [u, v] : edges) {
if (L.time[u] > L.time[v]) swap(u, v);
cout << ans[v] << ' ';
}
cout << '\n';
}
}
