AtCoder Regular Contest 118 F - Growth Rate

問題

atcoder.jp

解法

数列 $X$ を数える代わりに整数列 $D = (D _ 1, \ldots, D _ {N + 1})$ を $D _ 1 \coloneqq X _ 1,\ D _ i \coloneqq X _ i - A _ {i - 1} X _ {i - 1}\ (i \geq 2)$ と定めて $D$ を数える。

$D$ が次の条件をすべて満たすことが、$X$ が条件を満たすための必要十分条件である。

  1. $\displaystyle \sum _ {i = 1} ^ {N + 1} D _ i \prod _ {j = i} ^ N A _ j \leq M $
  2. $D _ 1\geq 1$
  3. $D _ i \geq 0\ (i=2,3,\ldots,N+1)$

$D _ 1$ だけ範囲が異なるのが嫌なので改めて $D _ 1\gets D _ 1 - 1$ および $\displaystyle M\gets M - \prod _ {i = 1} ^ N A _ i$ とすれば、$D$ が次の条件をすべて満たすことが、$X$ が条件を満たすための必要十分条件である。

  1. $\displaystyle \sum _ {i = 1} ^ {N + 1} D _ i \prod _ {j = i} ^ N A _ j \leq M $
  2. $D _ 1\geq 0\ (i=1,2,\ldots,N+1)$

$\displaystyle P _ i\coloneqq\prod _ {j = i} ^ N A _ j$ と定めれば、これは次のように書き直すことができる。

  1. $\displaystyle \sum _ {i = 1} ^ {N + 1} D _ i P _ i \leq M $
  2. $D _ 1\geq 0\ (i=1,2,\ldots,N+1)$

$D$ を桁 DP の要領で数える。より具体的には $\displaystyle D _ i = \sum _ {j = 0} ^ {i - 1} d _ {i, j} P _ {j + 1}$ および任意の $j=1,2,\ldots,i-1$ に対して $0\leq d _ {i, j} \lt A _ j$ を満たすように $D _ i$ を桁 $(d _ {i, 0} \ldots, d _ {i, i - 1})$ へと分解して考える。そして、$j$ 桁目 $d _ {j + 1,j},d _ {j + 2,j},\ldots, d _ {N + 1, j}$ を同時に決めることを $j$ の降順に行っていく。以降、場合分けが煩雑になることを避けるため、$i\leq j$ なる $i, j$ に対して $d _ {i, j} = 0$ とする。

さて、いま、$j \gt t\ (t \geq 0)$ なる $j$ に対する $d _ {\ast, j}$ を既に決めたとする。残りの $j\leq t$ に対する $d _ {\ast, j}$ が満たすべき条件は次のように書ける。

$$\sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ t d _ {i, j} \dfrac{P _ {j + 1}}{P _ {t + 1}} \leq \left\lfloor \dfrac{M - \sum _ {i = 1} ^ {N + 1} \sum _ {j = t + 1} ^ N d _ {i, j} P _ {j + 1}}{P _ {t + 1}} \right\rfloor$$

$0\leq d _ {i, j} \lt A _ j$ の条件より $\sum _ {i = 1} ^ {N + 1} \sum _ {j = t + 1} ^ N d _ {i, j} P _ {j + 1}\lt (N+1) P _ {t + 1}$ であるから、ある整数 $x\ (0\leq x\lt N + 1)$ が存在して上式は次のように表せる。

$$\sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ t d _ {i, j} \dfrac{P _ {j + 1}}{P _ {t + 1}} \leq \left\lfloor \dfrac{M}{P _ {t + 1}} \right\rfloor - x$$

以降 $\left\lfloor \dfrac{M}{P _ j} \right\rfloor$ という形が良く現れるので、これを $M _ j$ と定める。以上より、次のような動的計画法を考えることができる。

$$\begin{aligned}\mathsf{dp}(t, x) \coloneqq{}& \text{全ての $i\gt j \leq t$ なる $i, j$ に対する $d _ {i, j}$ の決め方であって、} \\ & \text{$\displaystyle \sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ t d _ {i, j} \dfrac{P _ {j + 1}}{P _ {t + 1}} \leq M _ {t + 1} - x$ であるようなものの個数}\end{aligned}$$

このとき、答えは $\mathsf{dp}(N, 0)$ である。

まず、$\mathsf{dp}(0, x)$ の計算を考える。これは $\displaystyle \sum _ {i = 1} ^ {N + 1} d _ {i, 0} \leq M _ 1 - x$ なる $d _ {1,0},\ldots,d _ {N+1,0}\geq 0$ の決め方の個数なので、$\displaystyle \mathsf{dp}(0, x) = \binom{N + 1 + M _ 1 - x}{N + 1}$ である。

続いて $t\gt 0,\ 0\leq x\leq N$ に対する $\mathsf{dp}(t, x)$ の計算を考える。$x$ として $0\leq x\leq N$ を満たすものだけを考えれば十分なのは先述の通りである。

$M _ t = \dfrac{M _ {t + 1} - (M _ {t + 1} \bmod A _ t)}{A _ t}$ より $R _ t \coloneqq M _ {t + 1} \bmod A _ t$ と定めると $M _ {t + 1} - x = A _ t M _ t + R _ t - x$ である。

また $\displaystyle \sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ t d _ {i, j} \dfrac{P _ {j + 1}}{P _ {t + 1}} = \sum _ {i = t + 1} ^ {N + 1} d _ {i, t} + A _ t \sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ {t - 1} d _ {i, j} \dfrac{P _ {j + 1}}{P _ t}$ であるから、条件 $\displaystyle \sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ t d _ {i, j} \dfrac{P _ {j + 1}}{P _ {t + 1}} \leq M _ {t + 1} - x$ は次のように整理できる。

$$\sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ {t - 1} d _ {i, j} \dfrac{P _ {j + 1}}{P _ t} \leq M _ t - \left\lceil\dfrac{\left(\sum _ {i = 1} ^ {N + 1} d _ {i, t}\right) + x - R _ t}{A _ t}\right\rceil.$$

$\displaystyle s\coloneqq \sum _ {i = t + 1} ^ {N + 1} d _ {i, t}$ と定めるとさらに次のように表せる。

$$\sum _ {i = 1} ^ {N + 1} \sum _ {j = 0} ^ {t - 1} d _ {i, j} \dfrac{P _ {j + 1}}{P _ t} \leq M _ t - \left\lceil\dfrac{s + x - R _ t}{A _ t}\right\rceil.$$

従って、$\displaystyle \sum _ {i = t + 1} ^ {N + 1} d _ {i, t} = s,\ 0\leq d _ {i, t}\lt A _ t$ なる整数列 $(d _ {t + 1, t}, d _ {t + 2, t},\ldots,d _ {N + 1, t})$ の個数を $f(s)$ とすると、次が成り立つ。

$$\mathsf{dp}(t, x) = \sum _ {s = 0} ^ {(N-t+1)(A _ t - 1)} f(s) \cdot \mathsf{dp}(t - 1, \lceil (s + x - R _ t) / A _ t \rceil).$$

$0\leq x\leq N,\ 0\leq s\leq (N - t + 1)(A _ t - 1),\ 0\leq R _ t\leq A _ t - 1$ の下で、確かに再び $0\leq \lceil (s + x - R _ t) / A _ t \rceil\leq N$ が成り立っていることに注意する。

さて、$0\leq y\leq N$ なる整数 $y$ に対して $y = \lceil (s + x - R _ t) / A _ t \rceil$ を満たす $s$ の範囲は $(y A _ t + R _ t - x) - A _ t\lt s\leq (y A _ t + R _ t - x)$ である。

従って、$f(l,r)\coloneqq \sum _ {s = l + 1} ^ r f(s)$ と定めると次が成り立つ。

$$\mathsf{dp}(t, x) = \sum _ {y = 0} ^ {N} f((y A _ t + R _ t - x) - A _ t, y A _ t + R _ t - x) \cdot \mathsf{dp}(t - 1, y).$$

区間 $((y A _ t + R _ t - x) - A _ t, y A _ t + R _ t - x\rbrack$ の長さはちょうど $A _ t$ なので、$f((y A _ t + R _ t - x) - A _ t, y A _ t + R _ t - x)$ は、$0$ 以上 $A _ t$ 未満の整数からなる長さ $N - t + 2$ の数列であって、総和が $y A _ t + R _ t - x$ であるようなものの個数に一致する。

この数え上げは有名問題で、包除原理を用いることで次を得る。

$$\begin{aligned}& f((y A _ t + R _ t - x) - A _ t, y A _ t + R _ t - x) \\ &{}=\sum _ {i = 0} ^ {N - t + 2} (-1) ^ i \binom{N - t + 2}{i} \binom{(y A _ t + R _ t - x) - i A _ t + (N - t + 1)}{N - t + 1}.\end{aligned}$$

式の簡単のため $u\coloneqq N - t + 1$ と定めると、$\mathsf{dp}(t, x)$ は次のように表せる。

$$\begin{aligned} \mathsf{dp}(t, x) &{}= \sum _ {y = 0} ^ {N} \sum _ {i = 0} ^ {u + 1} (-1) ^ i \binom{u + 1}{i} \binom{(y A _ t + R _ t - x) - i A _ t + u}{u} \cdot \mathsf{dp}(t - 1, y)\\ &{}=\sum _ {y = 0} ^ {N} \mathsf{dp}(t - 1, y) \sum _ {i = 0} ^ {u + 1} (-1) ^ i \binom{u + 1}{i} \binom{(R _ t - x + (y - i) A _ t) + u}{u} .\end{aligned}$$

さて、$i\gt y$ のとき $R _ t + (y - i) A _ t \lt 0$ より $\displaystyle \binom{(R _ t - x + (y - i) A _ t) + u}{u} = 0$ である。また、$i\gt u + 1$ のとき $\displaystyle\binom{u + 1}{i} = 0$ である。従って、上式における内側の和で $i$ が動く範囲は $0$ 以上 $y$ 以下の範囲としてよい。

つまり、多項式 $a(z) = \displaystyle \sum _ {i = 0} ^ N (-1) ^ i \binom{u + 1}{i} z ^ i$ および $b(z) = \displaystyle \sum _ {i = 0} ^ N \binom{(R _ t - x + i A _ t) + u}{u} z ^ i$ について、$\lbrack z ^ y\rbrack (ab) = \displaystyle \sum _ {i = 0} ^ {u + 1} (-1) ^ i \binom{u + 1}{i} \binom{(R _ t - x + (y - i) A _ t) + u}{u}$ が成り立つ。

計算量を考える。$a(z), b(z)$ は各 $t,x$ に対してそれぞれ $O(N)$ 時間で計算できる。$b(z)$ の計算についてもう少し説明すると、これは $\displaystyle \binom{(R _ t - x + i A _ t) + u}{u} = \binom{(R _ t - (x - 1) + i A _ t) + u}{u}\cdot \dfrac{(R _ t - x + i A _ t) + u}{(R _ t - (x - 1) + i A _ t) + 1}$ を用いて SWAG で積をスライドさせながら求めることでならし $O(N)$ 時間が達成される。

$a(z)$ と $b(z)$ の積の計算は FFT を用いれば $O(N\log N)$ 時間である。これを全ての $t, x$ について行うので、全体 $O(N ^ 3 \log N)$ 時間となる。

本問題では $N$ は最大 $1000$ 程度にもなるので、このままでは実行時間制限に間に合わせることは難しい。

そこで、$A _ i \neq 1$ なる $i$ が高々 $\lfloor \log _ 2 M\rfloor$ 個しかないことに注目する。

$A _ t = 1$ の場合を考えると、任意の $x=0,1,\ldots,N$ について $\mathsf{dp}(t, x) = \mathsf{dp}(t - 1, x)$ が成り立つことが上記の議論から分かる。従って、$A _ t = 1$ の場合の遷移は $O(N)$ 時間で可能であるから、全体の計算時間は $O(N ^ 2 \log N \log M)$ となる。

これでも実行時間制限に間に合わせることは非常に厳しいが、$a(z)$ の FFT 結果をメモ化するなどの定数倍高速化を行うことで正答を得ることができる。

実装

atcoder.jp