# ALGOCUP3 - Editorial

Tester: Sankalp Gupta

MEDIUM

# PREREQUISITES:

Math, Matrix Exponentiation

# PROBLEM:

There are K villages numbered from 1 to K in the countryside. You live in a village numbered 1. All villages are interconnected. The cost of traveling from one village to another is 1
coin. You have exactly N coins with you. Your task is to find in how many ways you can spend exactly N coins so that you start from your home village 1and return to your home village in the end.As the answer can be large, print it modulo 10^9+7.

# QUICK EXPLANATION:

We will find the O(n) DP solution to the problem and then we will optimize it using the Matrix Exponentiation Concept.

# EXPLANATION:

Let’s find an O(n) solution first to the given problem.
We can see that there are 2 types of villages: home village and other villages. All other villages are identical.
Let’s denote A(n) as the answer to return to home village after n steps starting from home village.
Let’s denote B(n) as the answer to return to other village after n steps starting from home village.

Initially for n==0, A(0)=1 and B(0)=0 because in zero steps we can’t go from home village to other village.
We can clearly say that transition A(n)=(k-1)*B(n-1) and B(n)=(k-2)*B(n-1)+A(n-1) holds.
As there are (k-1) different other villages from where you can return to home village.
And we can return to other village from (k-2) other villages and one home village.
To optimize above O(n) solution we will use Matrix exponentiation.

Now representing the above set of equation in term of matrix
P = \begin{bmatrix} 0 & K-1\\ 1 & K-2 \end{bmatrix}

A0 = \begin{bmatrix} 1\\ 0 \end{bmatrix}

R = \begin{bmatrix} A(n)\\ B(n) \end{bmatrix}

To get the answer for n steps
we can say (P^n)A0 = R.
First entry of R is our required answer.
Now all left is to calculate Matrix P raise to the power n in O(log n) time.
We can do so by using recursive relation
P^n = P^x* P^x (if n is even)
= P^x * P^x * P (if n is odd )

where x = n/2.

# ALTERNATE EXPLANATION:

We can obtain a simple formula for the above solution : ( (k-1)^n + (k-1)*(-1)^n )/k.

# SOLUTIONS:

Setter's Solution
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define mod 1000000007

void multiply(ll mat,ll temp){
ll x=(mat*temp + mat*temp)%mod;
ll y=(mat*temp + mat*temp)%mod;
ll z=(mat*temp + mat*temp)%mod;
ll w=(mat*temp + mat*temp)%mod;
mat=x;
mat=y;
mat=z;
mat=w;
}

void power(ll mat,int n,int k){
ll a={0,k-1,1,k-2};
if(n==0 || n==1)return;
power(mat,n/2,k);
multiply(mat,mat);
if(n&1)multiply(mat,a);
}

int main(){
int t; cin>>t;
while(t--){
int n,k; cin>>n>>k;
ll mat={0,k-1,1,k-2};
power(mat,n,k);
cout<<mat<<"\n";
}
}

Tester's Solution
#include<bits/stdc++.h>
#define ll long long int
#define ull unsigned long long int
#define vi vector<int>
#define vll vector<ll>
#define vvi vector<vi>
#define vvl vector<vll>
#define pb push_back
#define mp make_pair
#define all(v) v.begin(), v.end()
#define pii pair<int,int>
#define pll pair<ll,ll>
#define vpii vector<pii >
#define vpll vector<pll >
#define ff first
#define ss second
#define PI 3.14159265358979323846
#define fastio ios_base::sync_with_stdio(false) , cin.tie(NULL) ,cout.tie(NULL)
ll power(ll a,ll b){ ll f=1; while(b>0){ if(b&1) f*=a; a*=a; b>>=1;} return f; }
ll power(ll a,ll b,ll m){ ll f=1; while(b>0){ if(b&1) f=(f*a)%m; a=(a*a)%m; b>>=1;} return f;}
bool pp(int a,int b) {return a>b;}
using namespace std;

ll m=1e9+7;
ll t;

void mul(ll t,ll f){
ll x = ((f * t)%m + (f * t)%m)%m;
ll y = ((f * t)%m + (f * t)%m)%m;
ll z = ((f * t)%m + (f * t)%m)%m;
ll w = ((f * t)%m + (f * t)%m)%m;

f = x;
f = y;
f = z;
f = w;
}

int cal(int n){
ll f = {{1,0},{0,1}};
while(n>0){
if(n&1)
mul(t,f);

mul(t,t);
n>>=1;
}
return f;
}

void solve(){
ll n,k;
cin>>n>>k;
assert(n>0&&n<=1e9);
assert(k>1&&k<=1e9);
t = 0;
t = k-1;
t = 1;
t = k-2;

cout<<cal(n)<<"\n";

}

int main()
{
fastio;
int t;
cin>>t;
assert(t>0&&t<=1e5);
while(t--){
solve();
}

return 0;
}

3 Likes

Can you also give an explanation for alternate solution.

I think A0 is {1,0} instead of {0,1}

updated.

1 Like

We can obtain alternate solution using a GP sum, like think using permutation we have exactly (k-1) ways to visit in 1…(n-1) turns, but we can not visit city 1 in (n-1)th step, so we have an recursive relation like A(n)=(k-1)^(n-1)-A(n-1)
so if we expend this we will get a GP (Geometric progression) and using sum formula of GP we can reach till this expression (provided in solution).

2 Likes

Okkk thanks , I appreciate it

I’ve found this formula another way. The total amount of ways to spend n coins is (k-1)^n – after each travel there are k-1 towns to visit next since you cannot return to the same town. From those (k-1)^n routes, I observed that there is an identical amount of routes that end in cities 2…k (since all cities are identical); and exactly 1 more route that ends in city 1 if n is even or 1 less route if n is odd. From that it’s easy to obtain the final formula: ((k-1)^n-1)/n+1 for even n and ((k-1)^n+1)/n-1 for odd n.

I am not sure on how to prove that there are exactly 1 more or 1 less routes to city 1 than to any other city, but the formula holds