この解説では

• 比較的簡潔だが低速な \(\mathrm{O}(N^{1.5} \log N)\) 解法
• 高速だが発展的な知識を要求される \(\mathrm{O}(N \log N)\) 解法

の 2 つの解法を説明しています。

\(\mathrm{O}(N^{1.5} \log N)\) 解法

\(F(x)\) を「サイコロを \(1\) 回振った時に \(n\) マス進む確率」の母関数とします。すなわち

\[F(x) = \frac{1}{K} \sum_{i=1}^K x^i\]

です。このとき、サイコロを \(n\) 回振ってマス \(k\) \((0 \leq k \lt N)\) にいる確率は \([x^k] F^n\) と表せます。よって、「サイコロを \(n\) 回振ってゴールしていない確率」\((0 \leq n \leq N)\) を \(a_n\) と置くと

\[a_n = \sum_{k=0}^{N-1} [x^k] F^n\]

と表すことが出来ます。「ちょうど \(n\) 回振ってゴールする確率」は \(a_{n-1} - a_n\) と表せるので \(a_0, a_1, \dots, a_n\) が列挙できれば答えが分かります。
式をもう少し整理します。\(a_n\) は \(F^n\) の \(0\) 次から \(N-1\) 次までの係数を集めたものなので、少し工夫すると

\[ \begin{aligned} a_n &= \sum_{k=0}^{N-1} [x^k] F^n \\ &= [x^{N-1}] (1 + x + \dots + x^{N-1}) F^n \end{aligned} \]

という形で表すことができるのがわかります。

ここからは \(K\) が大きい場合と小さい場合で \(2\) 種類の解法を使い分けることにします。

\(K\) が大きい場合

\(a_n\) をもう少し変形すると、

\[ \begin{aligned} a_n &= [x^{N-1}] (1 + x + \dots + x^{N-1}) F^n \\ &= [x^{N-1}] (1+x+\dots+x^{N-1}+x^N+\dots) F^n \\ &= [x^{N-1}] \frac{F^n}{1-x} \end{aligned} \]

となります。この式に

\[F(x) = \frac{1}{K} \sum_{i=1}^K x^i = \frac{x(1-x^K)}{K(1-x)}\]

を代入すると

\[ \begin{aligned} a_n &= [x^{N-1}] \frac{x^n (1-x^K)^n}{K^n(1-x)^{n+1}} \\ &= \frac{1}{K^n} [x^{N-1-n}] (1-x^K)^n \times \frac{1}{(1-x)^{n+1}} \end{aligned} \]

となります。上式の \((1-x^K)^n\) の部分は非 \(0\) の係数が \(0\) 次、\(K\) 次、\(2K\) 次、\(\dots\) と \(K\) 個置きにしか存在せず、また \((1-x^K)^n\) と \(\frac{1}{(1-x)^{n+1}}\) の係数の取得は適当な階乗・階乗逆元の前計算の元 \(\mathrm{O}(1)\) で行えるので、\(a_n\) は \(\mathrm{O}(N/K)\) で計算することが出来ます。よって \(a_0, a_1, \dots, a_N\) の列挙は \(\mathrm{O}(N^2 / K)\) で行うことが出来ます。

\(K\) が小さい場合

\(a_n\) の式を再掲します。

\[a_n = [x^{N-1}] (1 + x + \dots + x^{N-1}) F^n\]

実は \(a_0, a_1, \dots, a_N\) は、 ABC281 Ex でも登場した分割統治をすれば全体で \(\mathrm{O}(NK \log^2 N)\) で計算することが出来ます。

分割統治の方法を説明するために、はじめに TLE する解法を載せます。下のコードは \(a_0, a_1, \dots, a_N\) を少し工夫した分割統治によって \(\mathrm{O}(N^2 \log N)\) で計算したものです。(詳細はコードを読んでください。)

#include <iostream>
#include <vector>
using namespace std;

#include "atcoder/convolution.hpp"
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;

int N, K;
vector<mint> f; // F = 1/K sum{i=1..K} x^i
vector<mint> a; // (a_0, a_1, ..., a_N)

// input  : l, r, g = (sum{i=0..N-1} x^i) F^l mod x^N
// output : F^{r-l} mod x^N
vector<mint> dc(int l, int r, vector<mint> g) {
  int b = g.size();
  if (l + 1 == r) {
    a[l] = g.back();
    return f;
  }
  int m = (l + r) / 2;
  auto p = dc(l, m, g);
  g = atcoder::convolution(g, p), g.resize(b);
  auto q = dc(m, r, g);
  auto res = atcoder::convolution(p, q);
  if ((int)res.size() > N) res.resize(N);
  return res;
}

int main() {
  cin >> N >> K;
  a.resize(N + 1);
  f.resize(K + 1);
  for (int i = 1; i <= K; i++) f[i] = mint{K}.inv();
  dc(0, N + 1, vector<mint>(N, 1));
  for (int i = 0; i < N; i++) cout << (a[i] - a[i + 1]).val() << "\n";
}

このコードが \(\mathrm{O}(N^2 \log N)\) 掛かる理由は、再帰が呼ばれるごとに長さ \(N\) の形式的冪級数 g に関する畳み込みが行われるからです。
この g に手を加えることで高速化します。dc(l, r, g) が呼ばれた時点での g に対しては、高々 \(r-l-1\) 回しか \(F\) を掛けないため、g の項のうち必要な情報は上位 \((r-l-1)K + 1\) 項のみです。よって g の低次の項は適宜捨てても問題ないことが分かります。そこで、

  int b = g.size();

の部分を

  int b = min<int>(g.size(), (r - l - 1) * K + 1);
  g.erase(begin(g), begin(g) + g.size() - b);

と書き換えてみます。(このように書き換えても答えが変わらないことは確認できます。) こうすることで計算量が \(\mathrm{O}(N K \log^2 N)\) に落ちています。この計算量は、 dc(l, r, g) 内部で登場する g および p, q の次数が \(\mathrm{O}(\min((r-l)K, N))\) で抑えられていることから導出できます。(注：計算量は正確には \(\mathrm{\Theta}\left(N K \log N \log \frac{N}{K}\right)\) と考えられます。)

なお、類題の出題例は maspy さんの記事 の「パターン 2」という項にいくつか載っているので参考にしてみてください。

以上より \(\mathrm{O}(N^2 / K)\) の解法と \(\mathrm{O}(N K \log^2 N)\) の解法の 2 つの解法を得ることが出来ました。よって、\(K = \mathrm{O}(\sqrt{N} / \log N)\) 付近を境界として解法を使い分けることで、この問題を \(\mathrm{O}(N^{1.5} \log N)\) で解くことが出来ます。

実装は畳み込みだけ持っていれば簡潔に書ける一方、実行速度は言語によっては制約内に収めるのが困難だと思われます。ACL の畳み込みを利用して C++ で自然な実装をすれば多少の余裕を持って AC できると筆者は考えています。

\(\mathrm{O}(N \log N)\) 解法

はじめに必要な基本知識をまとめます。

ニュートン法

\(G(f) = 0\) を満たす \(f(x)\) を \(n\) 次の係数まで計算したいとする。(\([x^0]f(x)\) はあらかじめわかっているとする。)
\(\hat{f} \equiv f \pmod{x^d}\) がを満たす \(\hat{f}\) が分かっている時に

\[f \equiv \hat{f} - \frac{G(\hat{f})}{G'(\hat{f})} \pmod{x^{2d}}\]

が成り立つので上式から \(f \bmod{x^{2d}}\) を計算できるため、\(f\) の係数の既知の部分の長さを倍にすることができる。よって、\([x^0]f \equiv f \pmod{x^1}\) から始めて上式を \(\mathrm{O}(\log n)\) 回反復させることで \(f(x)\) を \(n\) 次まで計算できる。

ラグランジュの反転公式

形式的冪級数 \(F(x), G(x)\) は逆関数の関係にあるとする。(すなわち \(F(G(x)) = G(F(x)) = x\) が成り立つ。) 今、\([x^0]F(x) = [x^0]G(x) = 0\) かつ \([x^1]F(x) \neq 0, [x^1]G(x) \neq 0\) が成り立つ時、次の式が成り立つ。

\[[x^n] F(x)^k = \frac{k}{n} [x^{n-k}] \left(\frac{x}{G(x)}\right)^n\]

特に \(k=1\) のときに

\[[x^n] F(x) = \frac{1}{n} [x^{n-1}] \left(\frac{x}{G(x)}\right)^n\]

が成り立つ。また、任意の形式的冪級数 \(H(x)\) に対して次式が成り立つ。

\[[x^n] H(F(x)) = \frac{1}{n} [x^{n-1}] H'(x) \left(\frac{x}{G(x)} \right)^n\]

証明はここでは割愛します。(ニュートン法に関しては ABC260Ex 解説の末尾 に簡単に書かれています。ラグランジュの反転公式に関しては ABC222H 解説 に \(k=1\) の場合の証明が載っていて、同様の方針で一般の場合も証明できます。)

さて、解法を説明します。(途中まで \(\mathrm{O}(N^{1.5} \log N)\) 解法と部分的に同様ですが再掲します。)
\(F(x)\) を「サイコロを \(1\) 回振った時に \(n\) マス進む確率」の母関数とします。すなわち

\[F(x) = \frac{1}{K} \sum_{i=1}^K x^i\]

です。このとき、サイコロを \(n\) 回振ってマス \(k\) \((0 \leq k \lt N)\) にいる確率は \([x^k] F^n\) と表せます。よって、「サイコロを \(n\) 回振ってゴールしていない確率」\((0 \leq n \leq N)\) を \(a_n\) と置くと

\[a_n = \sum_{k=0}^{N-1} [x^k] F^n\]

と表すことが出来ます。「ちょうど \(n\) 回振ってゴールする確率」は \(a_{n-1} - a_n\) と表せるので \(a_0, a_1, \dots, a_n\) が列挙できれば答えが分かります。
上式だと \(k=0\) の時に少し都合が悪いので修正します。明らかに \(a_0 = 1\) であり、また \(n \gt 0\) のとき \([x^0] F^n = 0\) です。よって \(n\) の範囲を \(1 \leq n \leq N\) とすれば総和を取る範囲を \(1 \leq k \leq N-1\) にしても問題ないので、次式を元に考えることにします。

\[a_n = \sum_{k=1}^{N-1} [x^k] F^n\]

上式に反転公式を適用します。\(F(x)\) の逆関数を \(G(x)\) とします。

\[[x^k] F(x)^n = \frac{n}{k} [x^{k-n}] \left(\frac{x}{G(x)}\right)^k\]

です。形式的ローラン級数 (負冪の項が有限個存在する冪級数, 形式的冪級数と同様に扱える) を持ち出して式変形すると

\[ \begin{aligned} a_n &= \sum_{k=1}^{N-1} \frac{n}{k} [x^{k-n}] \left(\frac{x}{G(x)}\right)^k \\ &= \sum_{k=1}^{N-1} \frac{n}{k} [x^{-n}] \left(\frac{1}{G(x)} \right)^k \\ &= n [x^{-n}] \sum_{k=1}^{N-1} \frac{1}{k} \left(\frac{1}{G(x)} \right)^k \end{aligned} \]

になります。(\(\frac{1}{G(x)}\) は \(x^{-1}+\)(形式的冪級数) という形をした冪級数であるのに注意してください。)よって

\[H(G(x)) = \sum_{k=1}^{N-1} \frac{1}{k} \left(\frac{1}{G(x)} \right)^k\]

とおくと、\(H(G(x))\) の \(-n\) 次の係数から \(-1\) 次の係数までが計算できればこの問題を解くことが出来ると分かりました。
問題を分解します。次の 2 つの計算を行うことが出来ればよいです。

• \(F(x)\) から \(G(x) \bmod{x^N}\) を計算する
• \(G(x) \bmod{x^N}\) から \(H(G(x))\) の下位 \(N\) 次を計算する

まず、\(G(x) \bmod{x^N}\) の計算はニュートン法を用いて \(\mathrm{O}(N \log N)\) で行うことが出来ます。詳しく説明すると、\(G(x)\) は

\[A(G) = F(G) - x = 0, [x^0]G = 0\]

を満たす式なので、\(\hat{G} = G \bmod{x^d}\) が与えられたときに \(A(\hat{G}) \bmod{x^{2d}}\) と \(A'(\hat{G}) \bmod{x^{2d}}\) が高速に計算できればニュートン法で \(G\) の高速な計算を行えます。\(A(\hat{G})\) は

\[A(\hat{G}) = F(\hat{G}) - x = \frac{\hat{G}^{K+1} - \hat{G}}{K(\hat{G} - 1)} - x\]

なので形式的冪級数の inv, pow を利用すれば高速に計算できます。また、\(A'(\hat{G})\) は連鎖律を利用すると高速に計算できます。\(A'(G)\) と \(A(G)\) の間には

\[ \begin{aligned} &\frac{d(A(G))}{dx} = \frac{d(A(G))}{dG} \frac{dG}{dx}\\ &\to A'(G) = \frac{\frac{d}{dx} A(G)}{\frac{d}{dx} G} \end{aligned} \]

という関係式が成り立つので、\(A(\hat{G})\) を \(\text{mod }x^{2d+1}\) で計算しておくことで \(A'(\hat{G}) \bmod{x^{2d}}\) も高速に計算することが出来ます。

\(H(G(x))\) の計算も連鎖律を利用すればよいです。

\[ \begin{aligned} \frac{d(H(G))}{dx} &= \frac{d(1/G)}{dx} \frac{d(H(G))}{d(1/G)} \\ &= \frac{-\frac{d}{dx}G}{G^2} \sum_{k=1}^{N-1} \left(\frac{1}{G} \right)^{k-1} \\ &= -\left(\frac{d}{dx} G\right) \sum_{k=2}^N \left(\frac{1}{G}\right)^k \end{aligned} \]

であり、右辺は \(x^N\) を掛けると形式的冪級数の inv, pow で計算できる形です。よって右辺を計算した後に両辺を積分すれば \(H(G)\) を \(\mathrm{O}(N \log N)\) で計算することが出来ます。
以上の方針により、この問題を \(\mathrm{O}(N \log N)\) で解くことが出来て、これは十分高速です。