AtCoder Regular Contest 182 C問題メモ

AtCoder Regular Contest 182

C - Sum of Number of Divisors of Product

• 長さが $1$ 以上 $N$ 以下で各要素が $1$ 以上 $M$ 以下である整数列を良い数列と呼びます。
• また、良い数列に対するスコアを、その整数列に含まれる要素の総積を $X$ としたときの、$X$ の正の約数の総数とします。
• 良い数列は $\displaystyle \sum_{k=1}^{N}M^k$ 個ありますが、それらすべてに対するスコアの総和を $998244353$ で割った余りを求めてください。

• $1 \leq N \leq 10^{18}$
• $1 \leq M \leq 16$
• 入力はすべて整数

積の和典型。
いつも解答見て「お前だったのか」って言ってる気がする。

$X$ の約数の個数は、$X$ を $p_1^{a_1}p_2^{a^2}...$ と素因数分解したときの $(a_1+1)(a_2+1)...$ で求められる。

$M$ が小さいので、考えられる素因数は $(2,3,5,7,11,13)$ の6つのみである。

元の問題は、以下の問題に言い換えられる。

• 以下の試行を考える
  • 素因数に対応した箱を最大6つ用意し、はじめに1つずつボールを入れておく
  • 以下を 1回以上 $N$ 回以下繰り返す
    • $1～M$ の数を1つ決め、その素因数に従って箱にボールを入れる
      • $2$ を選んだら $2$ の箱に1個入れる
      • $4$ を選んだら $2$ の箱に2個入れる
      • $12$ を選んだら $2$ の箱に2個、$3$ の箱に1個入れる
  • 最終的な各箱に入ったボールの個数の積が、$X$ の正の約数の個数である
• 全ての操作列に対する答えの総和を求めよ

これは、更に $X$ の正の約数の個数は「最終的な各箱に入ったボールから1つずつ選んで黒く塗る方法の個数」と言い換えられる。

さらに、試行の過程で「ボールを箱に入れるとき、同時にボールを塗る/塗らないも決めちゃう」としてもよい。

既に塗られたボールの入った箱の集合 $S$ をもってDPすれば、 最終的に全ての箱に黒いボールが1つずつ入っている状態が、$X$ の正の約数の個数と一致する。

で、このDPの遷移を考えると、「$S$ が同じ状態は、各箱にボールが何個入っていようが、その後の遷移が同じ」とわかる。
制約の大きい厄介な $N$ の影響の大部分が無視できる。

よって、以下のDPが成り立つ。

• $\mathrm{DP}[i,S]=i$ 番目の数字まで決めて、既に黒く塗られたボールの入った箱の集合が $S$ である場合の数

これだとまだ $O(N2^{\pi})$（$\pi$ は $M$ 以下の素因数の個数=6）かかるが、 このDPは ある $S_x$ から $S_y$ への遷移が $i$ に依らず一定なので、行列累乗で高速化できる。

ちょうど $N$ 回なら単に $S$ の状態数 $2^{\pi} \times 2^{\pi}$ の行列累乗でいいが、 今回は「$N$ 回以下」を求める必要がある。

これは、行・列に1段追加して、以下のようにすればよい。

$\tilde{A} = \left( \begin{array}{ccccc} A_{11} & A_{12} & \ldots & A_{12^{\pi}} & 0 \\ A_{21} & A_{22} & \ldots & A_{22^{\pi}} & 0 \\ \vdots & \vdots & \ddots & \vdots & \vdots \\ A_{*} & A_{*} & \ldots & A_{*} & 0 \\ A_{*} & A_{*} & \ldots & A_{2^{\pi}2^{\pi}} & 1 \\ 0 & 0 & 0 & 0 & 1 \end{array} \right)$

（全部の添字を書くと面倒ごちゃごちゃになるので、適当に省略）

これで、「途中で全て選ばれた状態 $S=2^{\pi}$ になったら、右端列に適宜加算される」という遷移がなされるようになる。

最終的に、$\tilde{A}^{N+1}$ の $1～2^{\pi}$ 行目の右端列の総和が答えとなる。
（$N$ 回ちょうどまでの結果も右端に加算するため、1回多く累乗する）

$O(8^{\pi}\log{N})$ となる。

Python3