PASSTHRU - Editorial

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';
	}
}

I solved it using small to large merging. Here’s the code -
https://www.codechef.com/viewsolution/82561792

I also did the same using small-to-large merging, but I couldn’t remove TLE.
Can anyone help me with this , Please
https://www.codechef.com/viewsolution/82571268

The complexity of your code is roughly N^2 and it would be in a graph like 1-2, 2-3, 3-4, 4-5, 5-6 … (N-1)-N when all the vertices are of distinct colors.

Basically the small to large merging isn’t implemented correctly. You are running a loop after merging which will increase the time complexity, optimize it to calculate the answer while merging in the first loop itself (in the merge function in your code).

CodeChef | Competitive Programming | Participate & Learn. Why am I getting TLE ? I have used small to large merging.

Can you please check why my code is getting TLE on this ?

Can you please explain your idea briefly.

During the contest, I came up with a fairly straightforward solution. Firstly, I constructed an Euler tour for the tree and then store it in an array, hence each subtree could be perceived as a subarray. Now for each subarray, we must count the number of distinct colors of its elements, which is a well-known data structures problem (For instance: this task right here)

Here is my implementation: CodeChef | Competitive Programming | Participate & Learn

UPD: I was finally able to get AC after removing segment tree after I realized it was an overkill. But I am still not able to figure out why is it giving TLE. Isn’t the time complexity O(nlg2n) ? And if it is, why doesn’t it pass ?

Although they have the same complexity, most of the time fenwick trees run much faster than segment trees IRL

Your code does have \mathcal{O}(N\log^2 N) complexity.

The intended complexity for this problem is \mathcal{O}(N\log N), so \mathcal{O}(N\log^2 N) would only pass if you have good constant factor (for example, using fenwick instead of segtree as mentioned above, or not using a data structure like that at all).

@iceknight1093
Any suggestions on my above-mentioned submission, and ways to remove TLE, please?

My suggestion is for you to look through why small-to-large merging works in the first place :slight_smile:

bhavaygoyal’s comment above illustrates the issue with your code perfectly: if you merge small-to-large, but then iterate across the entire set afterwards, it obviously defeats the point, no? The idea behind small-to-large is to not iterate across as many things as possible.

Basically, your code is \mathcal{O}(N^2) and no small constant-factor optimization is going to fix that, you need a couple more ideas.

Alternately, you could read the editorial which details a solution that doesn’t use small to large at all.

1 Like

Thanks a lot, I will definitely work on it!

Understood the concept of auxiliary Tree, but what is the significance of DFS at the end and why is it working?

Is it something related to pushing down the edge weight to vertex, and edge answer is equivalent to value accumulated at the vertex having higher DFS time?

The final DFS performs the analogue of prefix sums on the tree.

As noted in the solution, once you build the auxiliary trees you want to add some value to each edge in the auxiliary tree, which I break into adding that value to several paths between a node and its ancestor.

Then, we’d like to use the trick I mentioned to process subarray additions offline as mentioned here:

To generalize this to a tree you use subarray sums, and of course you need a DFS to compute that.

The idea is, if you want to add c to the path from u to v where u is an ancestor of v, you increase the value at v by c and decrease the value at u by c.
When you take subarray sums, you’ll notice that these two changes together contribute exactly c to every edge on the u\to v path, and 0 to every other edge, which is exactly what you want.