PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Heavy-Light decomposition, Segment Trees
PROBLEM:
You’re given a tree on N vertices. Some vertices are marked.
The score of the tree is defined to be the sum of \text{depth}(\text{lca}(u, v)) across all pairs (u, v) of marked vertices.
There are Q updates to the tree.
Each update gives you a vertex u, and you must toggle the state of all vertices in the subtree of u.
Compute the score of the tree after each update.
EXPLANATION:
The first thing we need to do is make the score computation more tractable; since quadratic is obviously terrible.
It makes sense to condition on the LCA rather than the pairs of marked vertices; since there are only N options for the LCA.
Let c_u denote the number of marked vertices present in the subtree of u.
Note that only some pairs among these c_u vertices can have their LCA be u.
Specifically, for each pair (x, y) of these c_u vertices,
- If there exists a child v of u such that the subtree of v contains both x and y, then \text{lca}(x, y) \ne u.
- Otherwise, \text{lca}(x, y) = u.
It’s now easy to see that the number of pairs of marked vertices whose LCA equals u, is exactly
where v ranges across all children of u.
That is, take all pairs, and then subtract out pairs that lie within the same child subtree.
So, the score of the tree then equals
Let’s look at the above expression a bit more carefully.
Specifically, let’s look at the coefficient of the term \binom{c_u}{2}.
- When looking at u itself, it appears with a coefficient of \text{depth}(u).
- When looking at \text{parent}(u), it appears with a coefficient of -\text{depth}(\text{parent}(u)).
However, note that \text{depth}(\text{parent}(u)) = \text{depth}(u) - 1.
So, the overall coefficient of \binom{c_u}{2} is just 1.
(Observe that this works for u = 1 as well, even though \text{parent}(u) doesn’t exist for it).
Thus, the score of the tree is simply
This is a much nicer expression to work with.
Next, we need to understand how to maintain this quantity across subtree flip updates.
Suppose we flip the subtree of v.
Looking at which values of c_u change,
- For each u in the subtree of v, c_u changes to (\text{sub}_u - c_u), where \text{sub}_u is the size of the subtree of u.
This is because marked and marked vertices flip for this subtree. - Further, for each u that’s an ancestor of v, c_u increases by (\text{sub}_v - 2c_v), because there are (\text{sub}_v - c_v) new marked nodes in the subtree, while c_v existing marked nodes become unmarked.
So, we need a data structure that preserves \sum \binom{c_u}{2}, while supporting path additions to c_u and also subtree “flip” operations.
Let’s work on each of those operations separately.
First, let’s look at just the path addition operations.
For simplicity, assume we just have the array c, and we want to add some integer K to the subarray [L, R] of c, while maintaining \sum \binom{c_u}{2}.
First, note that \binom{c_u}{2} = \frac{c_u \cdot (c_u-1)}{2}.
The division by 2 is a constant that’s present in every term, so we can ignore it for now and just divide by 2 in the end.
This means we want to know \sum c_u\cdot (c_u-1) = \sum\left(c_u^2 - c_u \right)
Now, maintaining this can in fact be done using a segment tree and a bit of algebra, because
So, if we know \sum (x_i^2 - x_i) for some range, and we want to add K to all the x_i in that range, the resulting value is
The change in value is hence \sum (2x_iK + K^2 - K), which can be easily computed knowing only \sum x_i and the length of the range.
This allows us to maintain the value we want while processing range additions, of course with lazy propagation.
As for path additions on a tree, there’s a standard method of converting range updates/queries on an array to path updates on a tree with an additional \mathcal{O}(\log N) factor: heavy-light decomposition.
This solves the path update part.
We also have to support subtree “flip” operations, involving replacing c_u with \text{sub}_u - c_u for all i in some subtree.
The simplest way to do this is to simply store all the information we need corresponding to c_u, and then store the same information corresponding to (\text{sub}_u - c_u) as well.
This means all the extra information needed by the segment tree to update the required answer, including lazy tags and such.
A “flip” operation then corresponds to just swapping all the pieces of information within a subtree; which after an Euler tour of the tree corresponds to swapping the information on some range.
This can again be handled with lazy propagation, since the swap and range addition updates play nice with each other - figuring out how to combine them is not too hard.
We’re thus able to process any single update in \mathcal{O}(\log^2 N) time with HLD while maintaining the answer, which is fast enough here.
TIME COMPLEXITY:
\mathcal{O}(N + Q\log^2 N) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
// store sum(x^2 - x) in each node
// also need sum of x, 1 in each node
// for flip, also store the same information for zeros
struct Node {
using T = array<ll, 3>;
T unit = T();
T f(T a, T b) {
T res = unit;
for (int i = 0; i < size(res); ++i) res[i] = a[i] + b[i];
return res;
}
Node *l = 0, *r = 0;
int lo, hi;
ll madd = 0;
bool mflip = false;
T val = unit, val2 = unit;
Node(int _lo,int _hi):lo(_lo),hi(_hi){ }
void init(int pos, ll x) {
if (pos < lo or pos >= hi) return;
if (lo+1 == hi) {
val[2] = val2[2] = 1;
val2[1] = x;
val2[0] = x*x - x;
}
else {
push();
l->init(pos, x), r->init(pos, x);
val = f(l->val, r->val);
val2 = f(l->val2, r->val2);
}
}
T query(int L, int R) {
if (R <= lo || hi <= L) return unit;
if (L <= lo && hi <= R) return val;
push();
return f(l->query(L, R), r->query(L, R));
}
void upd(auto &arr, ll d) {
// (x+d)^2 - (x+d) = x^2 - x + (2d)x + (d^2 - d)
arr[0] = arr[0] + 2*d*arr[1] + (d*d - d)*arr[2];
arr[1] = arr[1] + d*arr[2];
}
void add(int L, int R, ll x) {
if (R <= lo || hi <= L) return;
if (L <= lo && hi <= R) {
madd += x;
upd(val, x);
upd(val2, -x);
}
else {
push(), l->add(L, R, x), r->add(L, R, x);
val = f(l->val, r->val);
val2 = f(l->val2, r->val2);
}
}
void flip(int L, int R) {
if (R <= lo || hi <= L) return;
if (L <= lo && hi <= R) {
mflip ^= 1;
madd *= -1;
swap(val, val2);
}
else {
push(), l->flip(L, R), r->flip(L, R);
val = f(l->val, r->val);
val2 = f(l->val2, r->val2);
}
}
void push() {
if (!l) {
int mid = lo + (hi - lo)/2;
l = new Node(lo, mid); r = new Node(mid, hi);
}
if (mflip) {
l->flip(lo,hi);
r->flip(lo,hi);
mflip = 0;
}
if (madd) {
l->add(lo,hi,madd);
r->add(lo,hi,madd);
madd = 0;
}
}
};
// HLD<true> if values are on edges, remember to define segtree update/query properly!
template <bool VALS_EDGES> struct HLD {
int N, tim = 0;
vector<vector<int>> adj;
vector<int> par, siz, depth, rt, pos;
Node *seg;
HLD(vector<vector<int>> adj_)
: N(size(adj_)), adj(adj_), par(N, -1), siz(N, 1), depth(N),
rt(N),pos(N){ dfsSz(0); dfsHld(0); seg = new Node(0, N); }
void dfsSz(int v) {
if (par[v] != -1) adj[v].erase(find(begin(adj[v]), end(adj[v]), par[v]));
for (int& u : adj[v]) {
par[u] = v, depth[u] = depth[v] + 1;
dfsSz(u);
siz[v] += siz[u];
if (siz[u] > siz[adj[v][0]]) swap(u, adj[v][0]);
}
}
void dfsHld(int v) {
pos[v] = tim++;
for (int u : adj[v]) {
rt[u] = (u == adj[v][0] ? rt[v] : u);
dfsHld(u);
}
}
template <class B> void process(int u, int v, B op) {
for (; rt[u] != rt[v]; v = par[rt[v]]) {
if (depth[rt[u]] > depth[rt[v]]) swap(u, v);
op(pos[rt[v]], pos[v] + 1);
}
if (depth[u] > depth[v]) swap(u, v);
op(pos[u] + VALS_EDGES, pos[v] + 1);
}
void modifyPath(int u, int v, ll val) {
process(u, v, [&](int l, int r) { seg -> add(l, r, val); });
}
ll queryPath(int u, int v) { // Modify depending on problem
ll res = 0;
process(u, v, [&](int l, int r) {
res += seg -> query(l, r)[0];
});
return res;
}
ll querySubtree(int v) {
return seg -> query(pos[v] + VALS_EDGES, pos[v] + siz[v])[0];
}
void modifySubtree(int v) {
seg -> flip(pos[v] + VALS_EDGES, pos[v] + siz[v]);
}
};
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n, q; cin >> n >> q;
vector par(n, -1);
for (int i = 1; i < n; ++i) {
cin >> par[i];
--par[i];
}
vector adj(n, vector<int>());
for (int i = 1; i < n; ++i) {
adj[i].push_back(par[i]);
adj[par[i]].push_back(i);
}
HLD<false> H(adj);
for (int i = 0; i < n; ++i) {
H.seg -> init(H.pos[i], H.siz[i]);
}
while (q--) {
int x; cin >> x;
--x;
int pre = H.seg -> query(H.pos[x], H.pos[x]+1)[1];
H.modifySubtree(x);
int cur = H.seg -> query(H.pos[x], H.pos[x]+1)[1];
int dif = cur - pre;
int p = H.par[x];
if (p != -1) {
H.modifyPath(0, p, dif);
}
cout << H.querySubtree(0)/2 << ' ';
}
cout << '\n';
}
}