TREEWALK - Editorial

Problem Link :

Division 1

Division 2

Author : Jatin Yadav

Tester : Misha Chorniy

Editorialist : Anand Jaisingh

Pre-Requisites :

Matrices, Cayley Hamilton theorem, Polynomial Interpolation

Problem :

Given a tree consisting of N nodes rooted at node r, find the number of walks of length K between node r and all nodes of the tree.

Explanation :

I think finding the number of walks of length K in an arbitrary graph between 2 nodes i,j is a well known problem, and can be solved in O(N^{3} \cdot log_{2}(K) ) using matrix exponentiation. For example, you can read about it here

Basically, let matrix M be the adjacency matrix of a given graph G. For example, consider the tree given in the sample of the problem statement. It’s adjacency matrix is :

0 & 1 & 1 & 1 & 0 \newline
1 & 0 & 0 & 0 & 0 \newline
1 & 0 & 0 & 0 & 0 \newline
1 & 0 & 0 & 0 & 1 \newline
0 & 0 & 0 & 1 & 0

Now, the number of paths of length K between some fixed vertex i and all vertex j is given by the product M^{K} \cdot V, where the i^{th} component of column vector V of size N is 1, and all others are 0.
So, we can solve the 1^{st} sub task using this method, exponentiating the adjacency matrix.

For the second sub task, we find that its going to be very hard to find M^{K} in the given TL. To pass this sub task, we can use DP, where :


Dp[i][j]= \sum_{u} dp[u][j-1], for all u adjacent to i and j \ge 1 .

Now, coming to the 3^{rd} sub-task, we find that none of the above 2 methods are going to be sufficient. We need to do something more technical. Important Observation : The graph given in the input is always a tree, and not any arbitrary graph.

Also, one more important observation we can make is that in the last part of the matrix exponentiation, we are doing a lot of use-less multiplications, since the vector V consists of only 1 non-zero entry. Basically, only the elements of the r^{th} column of M^{K} are needed, since all other entries are always multiplied by 0.

Now, before proceeding further, you need to read about cayley hamilton theorem, from here and here

In-fact, we can use this to solve the problem. A theorem, Cayley Hamilton theorem can be used to find a single row or column of this matrix in O(N^2) time, but only if we can find its characteristic polynomial in such time too. Something that we need to be aware of is that if we can find the characteristic polynomial of a matrix, then we can express the powers >= N of a matrix in terms of the first N-1 powers of the matrix. This is because of the fact that each matrix satisfies it’s own characteristic polynomial.

However, the given blog link above deals only with the case of one dimensional linear recurrence. Our Recurrence is two dimensional, and until now, we have not found a method to find the characteristic polynomial.

Now, you need to read about polynomial interpolation before proceeding further.

Now, there are research based techniques based on finding the characteristic polynomial of a tree, it presents a neat and detailed way to find what we required. It is going to be useless for me to mention all those details all over again for me when I myself read it there, so you can access 2 of them here and here

Now, once we have found the characteristic polynomial, things should be easy to do. We can find the linear combination of I_{N},M,M^2,...M^{N-1} that equals M^K. Let this representation be [c_0,c_1,...c_{N-1}]. We can do this step in O(N^2 log_2 {K}) time.

So, M^{K}= \sum_{i=0}^{N-1} c_i \cdot M^i

However, we only need the r^{th} column of M^{K}, so , if we can calculate the r^{th} columns of I, M , M^2 ...M^{N-1}, then we’re good to go, we can just add the columns with their weight, i.e c_i for the i^{th} matrix. We can find the r^{th} column of all M,M^2,...M^{N-1} in O(N^2), but how ?

First, we know the r^{th} column of I_{N}. Assume we know the the r^{th} column of M^i, let’s prove we can find the r^{th} column of M^{i+1} in O(N) by induction :

M^{i+1}[x][r] = dot product of x^{th} row of M and r^{th} column of M^i. But, instead of calculating the dot product in a naive way, we can only go over the non-zero entries of the x^{th} row of M, For example :

M^{i+1}[x][r] = \sum_{u} M^{i}[u][r], for all u adjacent to x, since these are the only non-zero entries in row x of matrix M. To calculate all M^{i+1}[x][r], x=1,2,...N, this requires is O(N) operations. So to do this over all M,M^2,...M^{N-1}, we require O(N^2) time.

Now, the answer for node i is simply M^{K}[i][r]. If you need to familiarize yourself with cayley hamilton theorem first, then try this project euler problem.

Overall Complexity : O(N^2 \cdot log_{2}(K) )

Authors Solution : Link


I was able to solve it with an N^3 algorithm. I was able to reduce N by one half, because a tree is a bipartite graph, so I only needed to consider every other vertex starting from the root.

After bruteforcing the first n<=2N values for walks of length n from root to each vertex. O(N^2)
I used Gaussian elimination with partial pivoting to uncover the coefficients for the recurrence relation. I discover that the coefficients worked for every vertex.

One optimization is that the number of coefficients is < 3/4 of N’ (the number of vertices in the reduced graph). So, the O(N^3) Gaussian elimination algorithm uses 3/8 N. Also, the partial pivoting cuts down time reqts by another factor of 6. All together, the improvements cut the running time by a factor of 114.

The remainder of the algorithm was O(n^1.6 lg k) with Karatsuba multiplication.

How did you get the representation [c0,c1, …, c_N-1]?
I was trying to learn and implement the approach from this editorial, but didn’t get to finish it in time.

Actually I think it is possible to solve this problem even if the graph isn’t a tree, using Berlekamp–Massey algorithm to find the coefficients for the characteristic polynomial. BM is such a strange algorithm, from seeing data it creates the shortest linear feedback shift register describing the data, this time around the coefficients of the characteristic polynomial. So just feed it some data having the linear recursion relationship you want to find, and it will give you the coefficients.

I heard about BM when reading the editorial of a similar problem from sept long challenge, STCFOTT, from a solution by fjzzq2002, he has an explanation of the algorithm on codeforce explaining the main ideas.

My solution is actually much more brute than even using the characteristic polynomial. I simply simulate the process for a while, solving the small data set in O(N^2). And using BM I extrapolate the values from the simulation to solve the large data set. It’s such a strange and awesome algorithm.

The extrapolation runs in O(N^2+N \log(N) \log(K)) using fft or in O(N^2 \log(K)) without fft.


I found it is simpler to use the following publication: B. Mohar, Computing the characteristic polynomial of a tree to calculate characteristic polynomial. In essence:

if v is a leaf node:

f(v) = x;

f '(v) = 1;


f '(v) = f(v1) * f(v2)*… ; where v1,v2, … children nodes of v;


f(v) = f '(v) * (x-Sum { f '(vi)/f(vi) }); where sum is for all children nodes of v.

That leads to O(n^2) algorithm.

Can anyone plz explain me this step in the given Cayley Hamilton theorem link ?

Computing the first row of M2, M3, ..., Mk - 1 is not difficult either, and can be done is O (k^2) time.


Nice editorial! Just a small thing, I think your adjacency matrix is wrong in your example, it should be symmetric and have 0 on the diagonal.

1 Like

Made the change!

this is cool !

I added this link too, to the editorial.

Once you found the characteristic polynomial, so you get something like :

M^{n} = \sum_{i=0}^{N-1} c_i \cdot M^{i} .

Now, there’s a link in the editorial that shows how to get representation of M^{K} from this.

Doing this depends on the structure of a matrix. For example, if we have a matrix that corresponds to a 1-d recurrence of order k, we can notice that (k-1) columns of the matrix will have just 1 non-zero entry, and we don’t need to go over zero entries, when multiplying M^{i+1}=M \cdot M^{i}

In case the matrix is the adjacency matrix of a tree, I mentioned how to get one column. One can operate for a row similarly.

Got it.Thanks!!