以下の配列がわかっているとします。

\(dp[i][sum][rem]:=a^{i}\) 未満の整数 \(X\)で \(F(X)=sum\) かつ \(X\;mod\;p =rem\) を満たし、 \(s\) の条件も満たすものの個数 (\(K\) 以上のときは \(K\) とする)

この配列がわかっているとき \(i\) 番目に小さい、問題の条件を満たす \(X\) は上の桁から順にその桁の数字を決めることで求めることができます。

具体的には以下のように求めることができます。

\(num=i,sum=m,rem=t\) と初期化する。

\(j=n-1,n-2,...,0\) の順に以下の操作をする。

操作

\(s_{n-j}=1\) ならば \(k=0,1,2,...,\min(a-1,m)\) の順に \(dp[j][sum-k][rem-ka^{j}]\) と \(num\) の値を比較する。(\(rem\) の部分は \(mod\; p\) で考えてください。)

\(num\leq dp[j][sum-k][rem-ka^{j}]\) ならば \(X\) の \(j\) 桁目を \(k\) で確定させ、 \(sum\leftarrow sum-k,\;rem\leftarrow rem-ka^{j}\) と更新したのちに比較をやめる。

\(num>dp[j][sum-k][rem-ka^{j}]\) ならば \(num\leftarrow num-dp[j][sum-k][rem-ka^{j}]\) と更新して、比較を続ける。

この操作を高々 \(K\) 回することで答えが全て出力されます。

よって、答えを出力するパートでの計算量は \(O(Knm)\) で、全てのテストケースに対して十分高速です。

dpの求め方 (68 点)

\(s\) が全て \(1\) のときについて記述します。 \(0\) が出現したら適時スキップしてください。

\(i=0,1,...,n=1\) の順に、全ての \(sum=0,1,...,m,rem=0,1,...,p-1,k=0,1,...,\min(a-1,m)\) に対して、

\(dp[i+1][sum+k][rem+ka^{i}]\leftarrow dp[i+1][sum+k][rem+ka^{i}]+dp[i][sum][rem]\)

と更新する。

計算量は \(O(nm^{2}p)\) です。

dpの求め方 (100 点)

\(i=0,1,...,n-1\) の順に更新します。

\(68\) 点の式から

\(dp[i+1][sum][rem]=\sum_{k=0,1,...,\min(a-1,m)}dp[i][sum-k][rem-ka^{i}]\)

この式から以下のようになります。

\(dp[i+1][sum+1][rem+a^{i}-dp[i+1][sum][rem]=dp[i][sum+1][rem+a^{i}]-dp[i][sum-(1+\min(a-1,m))][rem-(1+\min(a-1,m))a^{i}]\)

よって、 \(r=0,1,...,p-1\) に対して、 \((rem, sum)=(r,0),(r+a^{i},1),(r+2a^{i},2),...\) の順番に差分更新することで \(dp\) を求めることができます。

計算量は \(O(nmp)\) です。

実装例 C++

ちなみにこれと同様のテクと累積和、二分探索を用いることで、計算量 \(O(nmp)\) の前計算をすることで答えを出力するパートの計算量を \(O(Kn\log(m))\) に落とすことができます。