Official

G - Sugoroku 5 Editorial by en_translator


This editorial two solutions:

  • a relatively simple but slow \(\mathrm{O}(N^{1.5} \log N)\) solution, and
  • a fast but advanced \(\mathrm{O}(N \log N)\) solution.

\(\mathrm{O}(N^{1.5} \log N)\) solution

Let \(F(x)\) be the generating function of the probability that you advance \(n\) squares in one roll, that is,

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

Then, the probability of being at square \(k\) \((0 \leq k \lt N)\) after \(n\) rolls is represented as \([x^k] F^n\). Thus, \(a_n\), the probability that you do not reach the goal after \(n\) rolls \((0 \leq n \leq N)\), can be expressed as

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

The probability of reaching the goal after exactly \(n\) rolls equals \(a_{n-1} - a_n\), so it is sufficient to enumerate \(a_0, a_1, \dots, a_n\) to find the answers.
Let us deform the expression a bit more. \(a_n\) is the sum of the coefficients of \(0\)-th through \((N-1)\)-th order of \(F^n\), so we can obtain the following form:

\[ \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} \]

Now, we try to use two solutions depending on whether \(K\) is large or small.

When \(K\) is large

\(a_n\) can be further deformed as:

\[ \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} \]

Assigning

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

to the expression above yields

\[ \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} \]

The \((1-x^K)^n\) part in the expression above has non-zero coefficients only on the \(0\)-th, \(K\)-th, \(2K\)-th \(\dots\) order terms, i.e. every \(K\) terms. Also, the coefficients of \((1-x^K)^n\) and \(\frac{1}{(1-x)^{n+1}}\) can be obtained in \(\mathrm{O}(1)\) time by appropriately precalculating factorials and their inverses. Thus, \(a_n\) can be found in \(\mathrm{O}(N/K)\) time. In total, \(a_0, a_1, \dots, a_N\) can be enumerated in a total of \(\mathrm{O}(N^2 / K)\) time.

When \(K\) is small

Take a look at the expression of \(a_n\) again.

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

In fact, \(a_0, a_1, \dots, a_N\) can be found in a total of \(\mathrm{O}(NK \log^2 N)\) time with divide-and-conquer, just as we did in ABC281 Ex.

In order to explain the divide-and-conquer, we first introduce a solution that results in TLE (time limit exceeded). The code below finds \(a_0, a_1, \dots, a_N\) in a total of \(\mathrm{O}(N^2 \log N)\) time with a bit clever divide-and-conquer technique. (For details, refer to the code.)

#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";
}

This code costs \(\mathrm{O}(N^2 \log N)\) time because formal power series g of length \(N\) is convolved every time the recursion occurs.
Let us tweak g to optimize it. When dc(l, r, g) is called, g is multiplied by \(F\) at most \(r-l-1\) times, so we only need to maintain the higher \((r-l-1)K + 1\) terms of g, and we may discard its lower order terms. Now let us replace

  int b = g.size();

with

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

(One can assert that this modification does not affect to the answer.) This way, the complexity is reduced to \(\mathrm{O}(N K \log^2 N)\), which is derived by the property that the degrees of g, p, and q appearing in dc(l, r, g) is bounded by \(\mathrm{O}(\min((r-l)K, N))\). (More precisely, the complexity is estimated as \(\mathrm{\Theta}\left(N K \log N \log \frac{N}{K}\right)\).)

Similar problems can be found in an article by maspy, section “パターン 2”. (Translator’s note: the article itself is written in Japanese, but you can find links to problems in CodeForces and AtCoder, where problem statements may be provided in English. Click the blue squares to navigate to them.)

Hence, we have obtained two solutions, one of \(\mathrm{O}(N^2 / K)\) and another of \(\mathrm{O}(N K \log^2 N)\). By picking \(K = \mathrm{O}(\sqrt{N} / \log N)\) as a boundary to choose which to use, the problem can be solved in a total of \(\mathrm{O}(N^{1.5} \log N)\) time.

Implementation requires only convolution, but it may be difficult to fit into the execution time limit in some languages. As long as you use convolution of ACL (AtCoder Library) and naturally implement with C++, we think you can get AC (accepted) with not too strict time limit.

\(\mathrm{O}(N \log N)\) solution

We first introduce prerequisites.

Newton’s method

Suppose we want to find the coefficients of up to \(n\)-th order of \(f(x)\) satisfying \(G(f) = 0\). When we know \(\hat{f}\) such that \(\hat{f} \equiv f \pmod{x^d}\), we have

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

which we can use to find \(f \bmod{x^{2d}}\), so we can double the length of the part for which coefficients are known. Starting from \([x^0]f \equiv f \pmod{x^1}\), we can repeat it \(\mathrm{O}(\log n)\) to obtain \(f(x)\) up to \(n\)-th order.

Lagrange inversion theorem

Suppose that formal power series \(F(x)\) and \(G(x)\) are inverse functions of each other (in other words, \(F(G(x)) = G(F(x)) = x\)). Now, if \([x^0]F(x) = [x^0]G(x) = 0\) and \([x^1]F(x) \neq 0, [x^1]G(x) \neq 0\), then we have:

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

Especially, when \(k=1\),

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

Moreover, for any formal power series \(H(x)\), we have:

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

We omit the proof here. (Newton’s method is briefly explained in the last part of the editorial of ABC260Ex. Lagrange inversion theorem is proven in the editorial of ABC222H for \(k=1\); general case can also be proven in the same idea.)

Now, we explain the solution. (The beginning is the same as the \(\mathrm{O}(N^{1.5} \log N)\) solution, but we will repeat it.)
Let \(F(x)\) be the generating function of the probability that you advance \(n\) squares in one roll, that is,

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

Then, the probability of being at square \(k\) \((0 \leq k \lt N)\) after \(n\) rolls is represented as \([x^k] F^n\). Thus, \(a_n\), the probability that you do not reach the goal after \(n\) rolls \((0 \leq n \leq N)\), can be expressed as

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

The probability of reaching the goal after exactly \(n\) rolls equals \(a_{n-1} - a_n\), so it is sufficient to enumerate \(a_0, a_1, \dots, a_n\) to find the answers.
This expression is not good for \(k=0\), so we will fix it. Obviously, \(a_0 = 1\), and \([x^0] F^n = 0\) for \(n \gt 0\). Thus, one can take the range of \(n\) to be \(1 \leq n \leq N\) to change the summation range to \(1 \leq k \leq N-1\). Now consider the following expression:

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

We apply the inversion theorem. Let \(G(x)\) be the inverse function of \(F(x)\), then

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

Using formal Laurent series (formal power series with a finite number of negative-power terms. It can be treated in the same manner as formal power series.), it can be deformed as follows:

\[ \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} \]

(Note that \(\frac{1}{G(x)}\) is a formal power series in the form of \(x^{-1}+\)(formal power series).) Thus, letting

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

it turns out that the problem can be solved by identifying the \((-n)\)-th through \((-1)\)-th coefficients of \(H(G(x))\).
Let us break down the problem. It is sufficient to compute the following two:

  • calculate \(G(x) \bmod{x^N}\) based on \(F(x)\).
  • calculate the lower \(N\) terms of \(H(G(x))\) based on \(G(x) \bmod{x^N}\).

First, \(G(x) \bmod{x^N}\) can be found in \(\mathrm{O}(N \log N)\) time using Newton’s method. Specifically, \(G(x)\) satisfies

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

so given \(\hat{G} = G \bmod{x^d}\), if one can find \(A(\hat{G}) \bmod{x^{2d}}\) and \(A'(\hat{G}) \bmod{x^{2d}}\) fast enough, then one can apply Newton’s method to compute \(G\) fast. \(A(\hat{G})\) satisfies

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

which can be computed using inverse and powering operations of formal power series. Also, \(A'(\hat{G})\) can be found using the chain rule. \(A'(G)\) and \(A(G)\) have the following relation:

\[ \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} \]

Thus, by computing \(\text{mod }x^{2d+1}\), we can also find \(A'(\hat{G}) \bmod{x^{2d}}\) too.

One can also apply chain rule to find \(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} \]

whose right hand side can be found using inverse and powering operations of formal power series. Thus, by finding the right hand side and then integrating, we can find \(\mathrm{O}(N \log N)\) in \(\mathrm{O}(N \log N)\) time.
Therefore, this problem can be solved in a total of \(\mathrm{O}(N \log N)\) time, which is fast enough.

posted:
last update: