Number P2106 Sam Rockwell Valley

Obviously matrix optimization problem. The only difficulty is the special judge k = 1when the answer is \ (10 \) instead of \ (9 \) .

#include <cstdio>
#include <cstring>

typedef unsigned int u32;
typedef unsigned long long ull;

const int MOD = 1e9 + 7;

class Matrix{
	private:
		u32 e[10][10];
	public:
		Matrix(){
			std::memset(e, 0, sizeof e);
		}
		u32 *operator[](const int& b){
			return e[b];
		}
		Matrix operator*(const Matrix& b)const{
			Matrix res;
			for(int i = 0; i < 10; ++i)
				for(int j = 0; j < 10; ++j)
					for(int k = 0; k < 10; ++k)
						res.e[i][k] = ((ull)res.e[i][k] + (ull)e[i][j] * b.e[j][k]) % MOD;
			return res;
		}
}ans, trans;

Matrix pow(Matrix a, ull b){
	Matrix res = a; --b;
	while(b){
		if(b & 1)
			res = res * a;
		a = a * a;
		b >>= 1;
	}
	return res;
}

ull k, rr;

int main(){
	std::scanf("%llu", &k);
	if(k == 1)
		return std::puts("10"), 0;
	for(int i = 0; i < 10; ++i)
		ans[1][i] = 1;
	for(int i = 0; i < 10; ++i)
		for(int j = 0; j < 10; ++j)
			if(i - j >= -2 && i - j <= 2)
				trans[i][j] = 1;
	if(k > 1)
		ans = ans * pow(trans, k - 1);
	for(int i = 1; i < 10; ++i)
		rr += ans[1][i];
	rr %= MOD;
	std::printf("%llu\n", rr);
	return 0;
}