PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Easy-Medium
PREREQUISITES:
DFS, binary search
PROBLEM:
You’re given a tree, each vertex is either black or white.
Define its score to be the sum of distances across all vertices with the same color.
You can choose a vertex and reattach its subtree to any of its ancestors.
Find the minimum possible score of the resulting tree.
EXPLANATION:
First, we need to write the score in a way that’s helpful to us.
This can be done by summing up the contribution of each edge.
Let’s define:
- N_0 and N_1 to be the total number of white/black vertices in the tree.
- S_{u, 0} and S_{u, 1} to be the number of white/black vertices in the subtree of u.
Then, for any u\ne 1, the number of same-color paths crossing the edge between u and its parent equals S_{u, 0} \cdot (N_0 - S_{u, 0}) + S_{u, 1}\cdot (N_1 - S_{u, 1}).
The score of the tree is obtained by summing this up across all u.
We now analyze how this changes with subtree movement.
Consider some vertex u, let p denote its parent and g denote its grandparent (i.e. the parent of p.)
Suppose we attach the subtree of u to g instead of p.
Then, observe that the only change occurs in the subtree of p - every other subtree has its white/black counts unchanged.
Analyzing the change in p, we lose S_{u, 0} white and S_{u, 1} black vertices.
So, the new contribution of the p-g edge (looking only at white) is
And its change compared to the original value is
The change in contribution for the black vertices is computed similarly.
In particular, if we define X_u = S_{u, 0}N_0 + S_{u, 0}^2 + S_{u, 1}N_1 + S_{u, 1}^2 to be a constant associated with u, the overall change of attaching u to g is
Now, note that attaching (the subtree of) u can be thought of repeatedly chaining the process of “attach u to its grandparent.”
So, if the ancestors of u are x_1, x_2, \ldots where x_1 is the parent of u; and we decide to attach u to x_k, then the overall change in cost will equal
Our goal is to choose u and k to minimize this sum.
To do that, observe that as i increases, both S_{x_i,0} and S_{x_i, 1} will increase.
So, the quantity \left(2S_{u, 0}S_{x_i, 0} + 2S_{u, 1}S_{x_i, 1} - X_u\right) is non-decreasing as we move towards the root.
Clearly, once this quantity becomes positive (or even 0, really) there’s no point moving any further upward.
So, the optimal solution is to move upwards as long as \left(2S_{u, 0}S_{x_i, 0} + 2S_{u, 1}S_{x_i, 1} - X_u\right) is negative, and stop once it isn’t.
Thus, our goal becomes to find the highest ancestor of u for which \left(2S_{u, 0}S_{x_i, 0} + 2S_{u, 1}S_{x_i, 1} - X_u\right) is negative; which can be done with binary search since this is monotonic.
Once the appropriate ancestor is known, it’s not too hard to compute the required change since it’s just a path sum on the tree.
In fact, by storing a list of ancestors when performing DFS, as well as suffix sums of their values S_{x_i, 0} and S_{x_i, 1} values, it’s possible to solve the problem without needing any special data structures at all.
Of course, if you want to you can always use binary lifting or something similar.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Editorialist's code (C++)
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
#include <cassert>
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector c(n, 0);
for (int &x : c) cin >> x;
vector adj(n, vector<int>());
for (int i = 0; i < n-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
}
vector<array<int, 2>> sub(n);
auto dfs = [&] (const auto &self, int u, int p) -> void {
++sub[u][c[u]];
for (int v : adj[u]) if (v != p) {
self(self, v, u);
sub[u][0] += sub[v][0];
sub[u][1] += sub[v][1];
}
};
dfs(dfs, 0, -1);
vector val(n, 0ll);
ll base = 0;
for (int u = 1; u < n; ++u) {
for (int i : {0, 1}) {
base += 1ll*sub[u][i] * (sub[0][i] - sub[u][i]);
val[u] += 1ll*sub[u][i] * (sub[u][i] + sub[0][i]);
}
}
ll best = 0;
vector<array<ll, 4>> path;
auto solve = [&] (const auto &self, int u, int p) -> void {
if (p != -1) {
if (path.empty()) path.push_back({sub[p][0], sub[p][1], sub[p][0], sub[p][1]});
else {
auto [x, y, px, py] = path.back();
path.push_back({sub[p][0], sub[p][1], px+sub[p][0], py+sub[p][1]});
}
}
if (p > 0) {
int lo = 0, hi = size(path);
while (lo < hi) {
int mid = (lo + hi)/2;
auto [x, y, px, py] = path[mid];
if (2ll*sub[u][0]*x + 2ll*sub[u][1]*y - val[u] < 0) hi = mid;
else lo = mid+1;
}
if (lo < size(path)) {
auto [x, y, px, py] = path.back();
if (lo > 0) {
auto [x2, y2, px2, py2] = path[lo-1];
px -= px2;
py -= py2;
}
best = min(best, 2ll*sub[u][0]*px + 2ll*sub[u][1]*py - val[u]*(ll(size(path)) - lo));
}
}
for (int v : adj[u]) if (v != p) {
self(self, v, u);
}
if (path.size()) path.pop_back();
};
solve(solve, 0, -1);
cout << base + best << '\n';
}
}