TREEPROD - Editorial



Panel Members

Problem Setter: Amit Pandey
Problem Tester: Pushkar Mishra
Editorialist: Prateek Gupta
Contest Admin: Praveen Dhinwa and Pushkar Mishra
Russian Translator: Sergey Kulik
Mandarin Translator: Hu Zecong
Vietnamese Translator: VNOI Team
Language Verifier: Rahul Arora




Topics - Dynamic Programming, Number Theory, Trees.
Problems - Distance in Tree


Given a weighted undirected tree of N vertices. You need to count the pair of vertices (i,j) such that product of weights of edges in path from i to j is divisible by given number M.


Solution for Subtask 1:

This subtask states that M is prime. Since, M is prime i.e it has no other prime factors except itself, we can mark the edges red which have their weights divisible by M. So, the problem becomes finding the number of paths having atleast one red coloured edge in them. This is because, that red edge will always be divisible by given prime number M. It is now a standard combinatorics problem.

First Approach

Finding the number of paths having at least one red edge is equivalent to counting the total paths in the tree excluding those which have no red edge in them. Total number of paths in the tree is infact (N*(N-1))/2. Now, how to find paths having no red edge between them?
If we somehow remove the edges which are read from the tree and then any path which will now exist wont be having that red edge. If thats the case, we would then be having several connected components in a tree. Each component will have a total of (Pi * (Pi - 1))/2 paths and each one wont have any of the red edge in it.
If K can be taken as the number of connected components formed by removing red edges and Pi as the size of the ith component formed, then the number of paths having atleast one red edge can be written in the form of following equation.

Total\ number\ of\ Paths\ having\ atleast\ one\ red\ edge\ = (N*(N-1))/2 - \sum_{i=1}^{k}(P_i*(P_i-1))/2

For those wondering if there is any other way, it can also be solved by applying DP over a tree. Let us see, how!

Second Approach

We can have the colour of the edge as red in case its weight is divisible by M. Let dpRed[i] denotes the number of paths starting from vertex i and ending in any other vertex in subtree of i having atleast one of the edges coloured red. Similarly, dpNotRed[i] denotes the number of paths starting from vertex i and ending in any other vertex in subtree of i having no edge coloured red. Now, for each vertex x, answer will be dpRed[x] + dpRed[x] * dpRed[childx] + dpNotRed[x] * dpRed[childx] + dpRed[x] * dpNotRed[childx] for all childx which are just immediate children of x. Look at the below pseudo code to understand it better.

Pseudo Code

void dfs(int curr, int prev)
    dpRed[curr] = 0;
    for ( int i = 0; i < (int)v[curr].size(); i++ ) {
      if ( v[curr][i].first == prev ) continue;
	  dfs(v[curr][i].first, curr);
	  int Red,NotRed;
	  if ( v[curr][i].second ) {
		Red = dpRed[v[curr][i].first] + dpNotRed[v[curr][i].first] + 1;
		NotRed = 0;
	  else {
		Red = dpRed[v[curr][i].first];
		NotRed = dpNotRed[v[curr][i].first] + 1;
	  ans += (lli)dpRed[curr]*(lli)NotRed;
	  ans += (lli)dpNotRed[curr]*(lli)Red;
	  ans += (lli)dpRed[curr]*(lli)Red;
	  dpRed[curr] += Red;
	  dpNotRed[curr] += NotRed;
  ans += (lli)dpRed[curr];

The above approach wont give correct results in case M is non-prime.

Solution for Subtask-2

This subtasks limits N i.e the number of vertices in the tree to 103.

Since, we now have a O(N^2) window now, we can run a breadth-first search from each vertex.
This breadth first search will explore all the paths starting from any vertex, lets call the vertex i.
So, if vertex j is not visited. It will assign the weight of the path ending at j as dis[j]. More formally, dis[j] = (dis[p]*edge\_Weight(j,p))%M where p is the vertex which has direct connection with j and lies in the path(i,j). We will run bfs for all such i vertices and add up the answer for all of them. Let us look at the below pseudo code to understand it better.

Pseudo Code

lli ans = 0;

for ( int i = 0; i < n; i++ ) {
	queue <int> q;
	vis[i] = i+1;
	dis[i] = 1;
	while ( !q.empty() ) {
		int f = q.front();
		for ( int j = 0; j < (int)v[f].size(); j++ ) {
			if ( vis[v[f][j].first] == i+1 ) continue;
			vis[v[f][j].first] = i+1;
			dis[v[f][j].first] = (dis[f]*v[f][j].second)%m;
	for ( int j = 0; j < n; j++ ) ans += (dis[j] == 0);

ans = ans/2;

But this will surely timeout in case N is little larger.

Solution for Subtask 3:

The catch to this subtask is the constraint on M which is 500. It can only have at max 4 distinct prime factors.


Lets take a path consisting of 3 vertices in which we have one edge from 1 <-> 2 having weight 3 and another edge from 2 <-> 3 having weight 5. Let’s take M as 30. Then, it as three distinct prime factors {2,3,5} which means our path will be represented as a frequency list of these primes. In our case , the product of the weights on path is 3^1 * 5^1 = 15. So, it will be represented as {0,1,1}. In case, the product was 45, then also it would have been represented as {0,1,1}. We won’t let any frequency in frequency list of given exceed the corresponding frequency for M which is {1,1,1} because it will anyway be divisible.

At each vertex, we keep the count of distinct frequency lists of paths starting from itself and ending at any of the vertex present in its subtree.

For e.g :-
Consider that M = 120.
Then prime_factors = {2, 3, 5}. Now take a tree of the following form :-

1 2 4

1 3 3

1 4 5

3 5 10

The HashMap with key-value pair at vertex 1 will be of the following form where key represents the frequency list and value represents the count of that particular frequency list.

[ {2, 0, 0}, 1), ({0, 1, 0}, 1), ({1, 1, 1}, 1), ({0, 0, 1}, 1} ]

Now, the main task is to count the number of paths at each vertex and add it up to the final answer.
so, for each vertex, we need to consider it’s every child such that one part of the path goes through that child and the other path goes through some other child in its subtree.

Let’s say all the paths contributed by a particular child be stored in map <vector <int>, int> child and all the paths contributed by parent without this particular child be sorted in map <vector <int>, int> parent.

Let’s count the total paths such that one part of the path goes through that child and the other path goes through some other child in its subtree.

Pseudo Code

void count_it(map< vector<int>, int> parent, map< vector<int>, int> child)

	map< vector<int>, int>::iterator it, iter;

	for(it = parent.begin(); it!= parent.end(); ++it)
		for(iter = child.begin(); iter!= child.end(); ++iter)
			bool check = true;
			for(int i = 0; i<prime_c.size();++i)
				if((it->first[i]) + (iter->first[i]) < prime_c[i])
					check = false;

				total_ans += 1ll*(it->second) * (iter->second);

Here, prime_f and prime_c denote the prime factors of M and theiry frequency list respectively.
So, we are iterating all pairs of paths in both the HashMaps and we are checking if multiplying them both would lead to higher powers compared to the powers of all primes in frequency list of M. If thats the case, we can take one such path from parent and one such path from child. So, total number of ways will be Comb(paths_from_parent, 1) * Comb(paths_from_child,1).

For details on the implementation have a look at the tester’s solution.


Overall complexity should be around O(N*P) which is enough to pass all kinds of subtasks stated in the problem. Here, constant P is atmost 500.



For the first subtask when M is prime, the problem can be solved in a simpler way.

Number of paths that have at least one red colored edge = \frac{n(n-1)}{2} - (number of paths that have no red colored edge)

To find the latter, simply delete the red colored edges and find the number of nodes in each resulting component. If Ai is the number of nodes in the ith component, then number of paths that have no red colored edge would simply be

\sum_{i=1}^{k}A_i * (A_i - 1)/2

for all k components.


I somehow forgot to mention this approach. Added it now. :slight_smile: