公式解説 では各遷移を毎回 \(O(K)\) で行っていますが、ビット演算による並列処理を使うと各遷移を \(O(1)\) で行うことができます。 本問では \(P\le 5\) なので \(0\) 以上 \(P\) 以下の整数は \(3\) ビットで表現できますが、さらに判定用に上位ビットを \(1\) ビット追加して、 \(4\) ビットずつ使うことにします。

「繰り上がり」を無視すれば、合計した状態は整数をそのまま合計すれば求まりそうですが、 \(P\) 以下に抑える処理をしないとオーバーフローしてしまいます。変数 x を \(4\) ビットごとに見て \(P\) を超えている場合に \(P\) に置き換えたものを表す関数を minP(x) とします。

この処理のため、 \(4\) ビットごとに次のようになっている変数を定義します。

• mask ： 0111 の繰り返し
• mmm ： 1000 の繰り返し
• ppp ： \(P\) の \(16\) 進法表現の繰り返し

結論としては minP(x) は、 x - ((((mmm + x - ppp) & mmm) * 7 >> 3) & (mmm + x - ppp)) で表せます。 これは (mmm + x - ppp) & mmm の部分を考えると、 \(x\) を \(4\) ビットごとに見たときに \(P\) 以上になっている箇所のみにビットが立っていることなどから分かります。

AC コード （PyPy3）