Codeforces Round #291 Div. 2 E - Darth Vader and Tree

Problem Statement:
514E - Darth Vader and Tree

Solution:
Learnt something new! A very clever idea to use fast matrix exponentiation to solve for the DP states. Check the official editorial for a better understanding :) (Thank you for the enlightenment!)

Let V(m) be the number of vertices that are located exactly m unit distance away from root. Hence we can have the following relationship: \(V(m) = \sum_{i=0}^{100} c[i] V(m-i) \), where c[i] is the number of edges with length i amongst the n edges. What we want to find is \(S(x) = \sum_{i=0}^{x} V(x) \). The clever idea is to transform these equations into matrix multiplications!



Let A = \( \left[
\begin{array}{cccccc} V(1) & V(2) & V(3) & \ldots & V(100) & S(100) \end{array} \right]
\)

and let B = \(
\left[
\begin{array}{cccccccc}

0 & 0 & 0 & 0 & \ldots & 0 &c[100] & c[100] \\
1 & 0 & 0 & 0 & \ldots & 0 & c[99] & c[99] \\
0 & 1 & 0 & 0 & \ldots & 0 & c[98] & c[98] \\
0 & 0 & 1 & 0 & \ldots & 0 & c[97] & c[97] \\
\vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots \\
0 & 0 & 0 & 0 & \ldots & 1 & c[1] & c[1] \\
0 & 0 & 0 & 0 & \ldots & 0 & 0 & 1 \\
\end{array}
\right]
\)

For each multiplication with B, the terms in A will increase by one, i.e. V(i) to V(i+1) and S(i) to S(i+1). Hence the problem is now equivalent to finding \(AB^{x-100}\), which can be done in logarithmic time using fast matrix exponentiation. Very clever construct :)

#include <iostream>
#include <cstdio>
#include <algorithm>
using namespace std;

long long MOD = (long long) 1e9 + 7LL;
int cnt[103];
long long A[103], B[103][103], C[103][103];
int N, X;
long long temp[103][103];


void fast_multiplication(int k) {
    if(k==1){
        for(int i=0;i<=100;++i){
            for(int j=0;j<=100;++j){
                C[i][j] = B[i][j];
            }
        }
        return;
    }
    if(k%2){
        fast_multiplication(k-1);
        for(int i=0;i<=100;++i){
            for(int j=0;j<=100;++j){
                temp[i][j] = 0;
                for(int k=0;k<=100;++k){
                    temp[i][j] += B[i][k] * C[k][j];
                    temp[i][j] %= MOD;
                }
            }
        }
        for(int i=0;i<=100;++i){
            for(int j=0;j<=100;++j){
                C[i][j] = temp[i][j];
            }
        }
    } else {
        fast_multiplication(k/2);
        for(int i=0;i<=100;++i){
            for(int j=0;j<=100;++j){
                temp[i][j] = 0;
                for(int k=0;k<=100;++k){
                    temp[i][j] += C[i][k] * C[k][j];
                    temp[i][j] %= MOD;
                }
            }
        }
        for(int i=0;i<=100;++i){
            for(int j=0;j<=100;++j){
                C[i][j] = temp[i][j];
            }
        }
    }
}

int main(){
    int d;
    scanf("%d%d",&N,&X);
    for(int i=0;i<N;++i){
        scanf("%d",&d);
        cnt[d]++;
    }
    for(int i=0;i<99;++i){
        B[i+1][i] = 1;
    }
    for(int i=0;i<100;++i){
        B[i][99] = cnt[100-i];
        B[i][100] = cnt[100-i];
    }
    B[100][100] = 1;
    A[0] = 1;
    for(int i=1;i<=100;++i){
        A[i] = 0;
        for(int j=1;j<=i;++j){
            A[i] += 1LL * cnt[j] * A[i-j];
            A[i] %= MOD;
        }
    } 
    if(X <= 100){
        long long ans = 0;
        for(int i=0;i<=X;++i){
            ans += A[i];
            ans %= MOD;
        }
        cout << ans << endl;
    } else {
        for(int i=0;i<=100;++i){
            A[101] += A[i];
            A[101] %= MOD;
        }
        fast_multiplication(X-100);
        long long ans = 0;
        for(int i=0;i<=100;++i){
            ans += C[i][100] * A[i+1];
            ans %= MOD;
        }
        cout << ans << endl;
    }
    return 0;
}