AtCoder Beginner Contest 357, D : 88888888

問題文

　https://atcoder.jp/contests/abc357/tasks/abc357_d

問題概要

　正整数をそれを十進数で表現する文字列に変換する関数を $f_s$ とし，正整数を十進で表す文字列を正整数に変換する関数を $f_i$ とする．また，文字列 $s$ を $k$ 回繰り返した文字列を $s^k$ で表す．
　正整数 $n$ に対し，$V_n$ を
\begin{equation*}
V_n := f_i( f_s( n )^n )
\end{equation*}
で定義する．
　$n$ が与えられるので，$V_n \bmod { 998{,}244{,}353 }$ を求めよ．

制約

• $1 \leq n \leq 10^{18}$

解法

　$\bmod$ をとることについては，演算の度に適切に処理するなり AtCoder Library (ACL) を使うなりすればよく，本質的な話でもないので $V_n$ 自体を求める体で書きます．
　$f_s( n )$ を一つずつ連結して，それらに対応する $f_i$ の値を求めることを考えます．右から連結していくと思うと連結の度に桁数分「シフト」するので，$f_i( f_s( n )^k )$ を $R_k$ と書くことにすれば，$n$ の桁数を $d$ として*1
\begin{equation*}
R_k = \begin{cases}
0 & \text{($k = 0$)} \\
R_{ k - 1 } \times 10^d + n & \text{(otherwise)}
\end{cases}
\end{equation*}
という漸化式が得られます．これを元に DP*2 をすれば $R_n$ ($= V_n$) を求めることはできますが，$\Omega( n )$ 時間かかって TLE します．
　上記の DP を高速化することを考えます．やや天下り的ですが，上記の漸化式の再帰ケースを
\begin{equation*}
R_k = n \times 1 + 10^d \times R_{ k - 1 }
\end{equation*}
と書いてみると，ふたつのベクトル $\begin{pmatrix} n \\ 10^d \end{pmatrix}$ と $\begin{pmatrix} 1 \\ R_{ k - 1 } \end{pmatrix}$ の内積になっているように見えます．またまた天下り的ですが，これを元に行列とベクトルの積で
\begin{equation*}
\begin{pmatrix}
1 & 0 \\
n & 10^d
\end{pmatrix}
\begin{pmatrix}
1 \\
R_k
\end{pmatrix}
=
\begin{pmatrix}
1 \\
R_{ k + 1 }
\end{pmatrix}
\end{equation*}
という式が作れます．行列をかけることで $R$ の添字が $1$ 進むので，$R_n$ を求めるには
\begin{equation*}
\begin{pmatrix}
1 & 0 \\
n & 10^d
\end{pmatrix}^n
\begin{pmatrix}
1 \\
R_0
\end{pmatrix}
=
\begin{pmatrix}
1 \\
R_n
\end{pmatrix}
\end{equation*}
を計算すればよいです．ここで，行列積は結合的な演算なので，行列の $n$ 乗の方を先に計算することができて，$n$ 乗は反復二乗法を用いることで乗算の回数を $\Theta( \log n )$ 回にすることができます．こうして行列の $n$ 乗を計算すれば，（かけられるベクトルが $\begin{pmatrix} 1 \\ 0 \end{pmatrix}$ なことから）結果の行列の $2$ 行 $1$ 列の要素が答えになっています．こういう方法を「行列累乗」と呼んだりします．
　計算量についてですが，行列のサイズが定数なので一回の乗算は $O( 1 )$ 時間で，全体では $\Theta( \log n )$ 時間となります．

コード

#include <atcoder/modint>
using namespace atcoder;
using mint = modint998244353;

using V = vector< mint >;
using M = vector< V >;

M mat_mul( const M &A, const M &B )
{
	M res = {
		{ 0, 0 },
		{ 0, 0 }
	};

	REP( i, 2 )
	{
		REP( j, 2 )
		{
			REP( k, 2 )
			{
				res[i][j] += A[i][k] * B[k][j];
			}
		}
	}

	return res;
}

M mat_pow( M &&A, LL N )
{
	M res = {
		{ 1, 0 },
		{ 0, 1 }
	};

	for ( ; N; N /= 2 )
	{
		if ( N & 1 )
		{
			res = mat_mul( res, A );
		}
		A = mat_mul( A, A );
	}

	return res;
}

int main()
{
	IN( LL, N );
	const int D = SZ( toString( N ) );

	const M A = mat_pow( {
			{ 1, 0 },
			{ N, mint( 10 ).pow( D ) }
			}, 
			N );

	cout << A[1][0].val() << endl;

	return 0;
}