H - Distinct Adjacent Editorial by MMNMM

閉じた式の導出

答えは \((-1) ^ N(M-1)+(M-1) ^ N\) です。 これを計算することで \(\Theta(N)\) 時間や \(\Theta(\log N)\) 時間でこの問題を解くことができます。

いくつかの方法でこれを導出することができます。

1. 包除原理

公式解説の包除原理を使うことで、求める場合の数は \[\sum _ {k=0} ^ N(-1)^k {} _ N\mathrm{C} _ kM ^ {\max\{1,N-k\}}\] とわかります。 これを整理すると、 \[\begin{aligned} \sum _ {k=0} ^ N(-1)^k {} _ N\mathrm{C} _ kM ^ {\max\{1,N-k\}}&=(-1) ^ N(M-1)+\sum _ {k=0} ^ N(-1)^k {} _ N\mathrm{C} _ kM ^ {N-k}\\ &=(-1) ^ N(M-1)+(M-1) ^ N \end{aligned}\] となります。

2. DP の高速化

公式解説の DP は、遷移が \(i\) によらないため、答えを次のような行列累乗の形で書くことができます。

\[a _ N=\begin{pmatrix}1&0\end{pmatrix}\begin{pmatrix}M-2&M-1\\1&0\end{pmatrix} ^ {N-1}\begin{pmatrix}0\\M\end{pmatrix}\]

これを整理すると、\(a _ N\) についての \(2\) 項間漸化式 \[\begin{aligned} a _ 2&=M(M-1)\\ a _ 3&=M(M-1)(M-2)\\ a _ n&=(M-2)a _ {n-1}+(M-1)a _ {n-2}\ (n\geq4) \end{aligned}\] が得られます。これを解くことで \(a _ n=(-1) ^ n(M-1)+(M-1) ^ n\) となります(遷移行列を対角化してもよいでしょう)。

実装例は以下のようになります。

#include <bits/extc++.h>
#include <atcoder/modint>

int main() {
    using namespace std;
    int N, M;
    cin >> N >> M;
    atcoder::static_modint<998244353> m = M - 1;
    cout << (m.pow(N) + (N & 1 ? -m : m)).val() << endl;
    return 0;
}
N, M = map(int, input().split())
p = 998244353

print((pow(M - 1, N, p) + pow(-1, N, p) * (M - 1)) % p)

posted:
last update: