Construct the Array [HackerRank 101 Hack 52 B]

https://www.hackerrank.com/contests/101hack52/challenges/construct-the-array

N個の数列を構成する。
各要素には1～Kを入れられる。
最初の要素は1, 最後の要素はXである。
他の要素を隣り合う要素が同じ数にならないように埋めていく。
何通りの埋め方があるか。

解法

https://www.hackerrank.com/contests/101hack52/challenges/construct-the-array/submissions/code/1304336933

部分点から考えていくとよい。
部分点では「dp[i][j] := i番目の要素まで正しく埋まっていて最後の数がjである組合せ」をやる。
これでは更新コストを無視してもO(NK)となるので完答できない。

これを最適化していくのだが、「jで全ての数を考えなくても=Xであるかどうかだけわかっていればいい」という考察を使う。
「dp[i][j] := i番目の要素まで正しく埋まっていて、最後の数がj=0なら!=Xでj=1なら=Xである組合せ」
これを更新する。
　
更新式を少しだけ解説する。
mint(K-2)でかけているのは、最後に選んだ数とXを抜いてK-2通りの遷移がある。
mint(K-1)でかけているのは、Xだけ抜いてK-1通りの遷移がある。

int N, K, X;
mint dp[101010][2];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K >> X;
    N -= 2;

    if (1 != X) dp[0][0] = 1;
    else dp[0][1] = 1;

    rep(i, 0, N) {
        dp[i + 1][0] = dp[i][0] * mint(K - 2) + dp[i][1] * mint(K - 1);
        dp[i + 1][1] = dp[i][0];
    }

    cout << dp[N][0] << endl;
}