PROBLEM LINK
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
DIFFICULTY:
Hard
PREREQUISITES:
Topics - Dynamic Programming, Number Theory, Trees.
Problems - Distance in Tree
PROBLEM:
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.
EXPLANATION:
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.
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.
Approach
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;
q.push(i);
vis[i] = i+1;
dis[i] = 1;
while ( !q.empty() ) {
int f = q.front();
q.pop();
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;
q.push(v[f][j].first);
}
}
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.
Approach
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 :-
5
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;
break;
}
}
if(check)
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.
COMPLEXITY
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.
