桁dpを代数的にやっつける

この記事はデータ構造とアルゴリズム Advent Calendar 2021 11日目の記事です。

はじめに

桁dpというのは動的計画法の一種である。例えば 「$N$ 以下の非負整数で10進数表現に4と9を含まないものはいくつあるか？」や「$N$ 以下の非負整数で10進数表現の桁和が $M$ であるものはいくつあるか？」のような問題は桁dp によって演算回数 $O(\log N)$ で解くことができる。本稿では桁dp を代数的に解釈する方法を解説する。あまり桁dp に慣れていない人はよい導入になるかもしれない。既に桁dp を深く理解している人にとっては、新しく解ける問題が増える訳ではないが、機械的に見通しよく問題を解けるようになるかもしれない。
なお、桁数の対数時間で解けるタイプの問題(「$N$ 以下」という制約がなく、桁数だけが指定されているもの)に対するアルゴリズムを代数的に理解するためには「競プロ典型90問5日目」を参照するとよい。

デジタル形式的冪級数

可換半環 $\mathbf{R}$ 上の形式的冪級数 $f\in\mathbf{R}[[x]]$ が整数 $m\ge 2$ について、$m$-デジタルであるとは、多項式の列 $(Q_k\in\mathbf{R}[x])_{k\ge 0}$ が存在し、$Q_k(0)=1$ ($\forall k\ge0$) と

f(x) = \prod_{k\ge0} Q_k\left(x^{m^k}\right)

を満たすことをいう。定義より $m$-デジタル冪級数は乗算について閉じていることが分かる。
本稿ではデジタル冪級数 $f$ について $[x^N]\,f(x)$ を計算するアルゴリズムとその適用例を紹介する。適用例によっては $Q_k(0)=1$ という条件を満たしていない場合(その代わりに積は有限に制限する)、非可換半環 $\mathbf{R}$ を考える場合、二変数のデジタル冪級数を考える場合などもある。しかし、簡単のために上記のようにデジタル冪級数を定義してアルゴリズムを考える。以下で説明するアルゴリズムは $f$ がデジタル冪級数の定義から多少はみ出した場合でも自明に一般化できる。

例

1/(1-x)

\frac1{1-x} = 1+x+x^2+x^3+\dotsb

は二進数展開を考えることで

(1+x)(1+x^2)(1+x^4)(1+x^8)\dotsm

と表現することができる。よって、任意の $k\ge0$ について $Q_k(x)=1+x$ とおくことで、$1/(1-x)$ は 2-デジタルであることが分かる。
より一般に $m$ 進数展開を考えると、$Q_k(x)=1+x+x^2+\dotsm+x^{m-1}$ とおくことで $1/(1-x)$ は $m$-デジタルであることが分かる。

また、 $m$-デジタルな形式的冪級数 $f(x)=\prod_{k\ge0} Q_k\left(x^{m^k}\right)$ について、$g(x):=f(x)/(1-x)$ を考えたい場合がよくある。このとき、上記のことから $ R_k(x):=(1+x+\dotsb+x^{m-1})Q_k(x)$ について $g(x)=\prod_{k\ge0} R_k\left(x^{m^k}\right)$ が成り立つ。

有理形式的冪級数

可換環上の任意の有理形式的冪級数 $f = P/Q$ は 2-デジタルである。ここで、 $P, Q\in\mathbf{R}[x]$ である。

\frac{P(x)}{Q(x)} = P(x)Q(-x)\frac1{Q(x)Q(-x)} = P(x)Q(-x)\frac1{Q_1(-x^2)}

ここで $Q(x)Q(-x)$ は偶多項式なので、ある多項式 $Q_1$ を用いて $Q_1(-x^2)$ と表すことができることを用いている。
$Q_0(x) = P(x)Q(-x)$, $Q_{1}(-x^2) = Q(x)Q(-x)$, $Q_{k}(-x^2)=Q_{k-1}(x)Q_{k-1}(-x),\quad k\ge 2$ とおくことで、2-デジタルであることが分かる。
詳しくは説明しないが、体上の任意の有理形式的冪級数は 任意の $m\ge2$ について $m$-デジタルである。

その他の例

1. $m$冪の整数を重複 $t$回まで許した分割の数の母関数 $f$ は $m$-デジタルである。任意の $k\ge0$ について $Q_k(x)=1+x+\dotsb+x^t$ とおけばよい。
2. 「$N$ が10進数表現に4と9を含まないならば1、そうでないときは0」という数列の母関数は 10-デジタルである。任意の $k\ge0$ について $Q_k(x)=1+x+x^2+x^3+x^5+x^6+x^7+x^8$ とおけばよい。
3. 「$N$ 以下の非負整数で10進数表現に4と9を含まないものの個数」という数列の母関数は 10-デジタルである。任意の $k\ge0$ について $Q_k(x)=(1+x+\dotsb+x^9)(1+x+x^2+x^3+x^5+x^6+x^7+x^8)$ とおけばよい。

アルゴリズム (LSB-first)

もし $Q_{k+1}$ が $Q_0,Q_1,\dotsc,Q_k$ から効率的に計算できるとき、$[x^N]\, f$ を効率的に計算するアルゴリズムを紹介する。任意の $m\ge2$ と $t\in\{0,1,\dotsc,m-1\}$ について、セクションオペレータ $\mathcal{S}_{t,\,m}$ を

\mathcal{S}_{t,\,m}\left(\sum_{n\ge0} a_n x^n\right) := \sum_{n\ge0} a_{nm+t}x^n

と定義する。すると、任意の形式的冪級数 $f\in\mathbf{R}[[x]]$ について

\begin{align*}
[x^N]\, f &= 
[x^N]\, \sum_{n\ge 0} f_n x^n\\
&= [x^N]\, \sum_{n\ge 0} f_{nm + (N\bmod m)} x^{nm + (N\bmod m)}\\
&= [x^{N-(N\bmod m)}]\, \sum_{n\ge 0} f_{nm + (N\bmod m)} x^{nm}\\
&= [x^{(N-(N\bmod m))/m}]\, \sum_{n\ge 0} f_{nm + (N\bmod m)} x^{n}\\
&= [x^{\lfloor N/m\rfloor}]\, \mathcal{S}_{N\bmod m,\, m}(f).
\end{align*}

が成り立つ。 任意の $m$-デジタル冪級数 $f\in\mathbf{R}[[x]]$ について

\begin{align*}
\mathcal{S}_{t,\,m}\left(\prod_{k\ge 0} Q_k\left(x^{m^k}\right)\right)
&=\mathcal{S}_{t,\,m}\left(Q_0(x)\right)\prod_{k\ge 1} Q_k\left(x^{m^{k-1}}\right)\\
&=\mathcal{S}_{t,\,m}\left(Q_0(x)\right)Q_1\left(x\right)\prod_{k\ge 1} Q_{k+1}\left(x^{m^{k}}\right)
\end{align*}

が成り立つ。このことから以下のアルゴリズムが得られる。

\begin{align*}
&\mathrm{Input}: Q_0,\ \mathrm{Output}: [x^N]\, \prod_{k\ge0}Q_k(x^{m^k})\\
&\\
&k \leftarrow 0\\
&P \leftarrow 1\\
&\mathbf{while}\ N > 0:\\
&\qquad  \mathrm{Compute}\ Q_k\\
&\qquad  P \leftarrow \mathcal{S}_{N \bmod m,\, m}(P Q_k)\\
&\qquad k \leftarrow k+1\\
&\qquad N \leftarrow \lfloor N/m\rfloor\\
&\mathbf{{return}}\ P(0)
\end{align*}

このアルゴリズムを有理形式的冪級数に適用したものがBostan–Moriのアルゴリズムに他ならない。
特に多項式 $P$, $Q$ について

[x^N]\, P(x)\prod_{k\ge0} Q(x^{m^k})

を計算するアルゴリズムは以下のようになる。

R.<x> = ZZ[]

# The section operator
def section(P, t, m):
  return R(P.list()[t::m])

# A function computing Nth term of P(x)Q(x)Q(x^m)Q(x^(m^2))Q(x^(m^3))...
def Nth_digital(N, P, Q, m):
  while N > 0:
    P = section(P * Q, N % m, m)
    N //= m
  return P(0)

print(Nth_digital(10^1000, R(1), 1+x+x^2+x^3+x^4, 2))

例えばある整数 $d$ について $Q$ の次数が $(m-1)d$ 以下で、 $P$ の次数が $d-1$ 以下のとき、このアルゴリズムの算術計算量は高々 $\mathsf{M}(d-1, (m-1)d)\lceil\log_m(N+1)\rceil$ である。ここで、$\mathsf{M}(n, m)$ は $\mathbf{R}$ 上の $n$ 次多項式と $m$ 次多項式の乗算の算術計算量である。

あまり重要ではない余談

ここで、$Q$ を固定して $P$ の更新に注目すると、線形な演算となっていることが分かる。よって、ある行列 $L_0,L_1,\dotsc,L_{m-1}$ が存在して、$N$ の $m$ 進数表現を $N_\ell\dotsm N_0$ とすると、 $P$ の更新は $L_{N_\ell}\dotsm L_{N_0} P$ と書ける。このとき、$N$ の $m$ 進数表現に登場する桁がすべて $h$ であった場合 $L_h^\ell P$ となる。これは算術計算量 $(\mathsf{MM}(d)+d^2)\log\lceil\log_m(N+1)\rceil$ で計算できる。ここで、$\mathsf{MM}(d)$ は $d\times d$ 行列の乗算の算術計算量である。
さらに、$L_h$ をベクトルに掛ける計算量は高々 $\mathsf{M}(d-1, (m-1)d)$ であることから、black-box linear algebra の手法を用いることで、小さな誤り確率を許せば算術計算量 $O(\mathsf{M}(d-1, (m-1)d) d + \mathsf{M}(d)\log\log N)$ で計算することもできる。

典型例

N以下の非負整数で10進数表現に4と9を含まないものの個数

ABC007 D, 禁止された数字
$\mathbf{R}=\mathbb{Z}$ とする。上述のように

Q(x)=\left(\sum_{n=0}^9 x^n\right)\left(\sum_{n\in\{0,1,\dotsc,9\}\setminus\{4,9\}} x^n\right)

について $[x^N]\, \prod_{k\ge0} Q(x^{10^k})$ が解である。
ABC007 Dのサンプルは以下のように計算できる。

print(10^18 + 1 - Nth_digital(10^18, R(1), (1+x+x^2+x^3+x^4+x^5+x^6+x^7+x^8+x^9)*(1+x+x^2+x^3+x^5+x^6+x^7+x^8), 10))
>>> 981985601490518016

算術計算量は $\mathsf{M}(1,17)\lceil\log_{10} (N+1)\rceil$ である。

N以下の非負整数で10進数表現の桁和がMであるものの個数

TDPC E, 数
$\mathbf{R}=\mathbb{Z}[y]/y^{M+1}$ とする。上記の議論と同様に

Q(x) = \left(\sum_{n=0}^9 x^n\right)\left(\sum_{n=0}^9 (yx)^n\right)

と定義すると、解は $[y^M]\,[x^N]\,\prod_{k\ge 0}Q(x^{10^k})$ である。

桁和を $D$ で割った余りが $M$ であるものの個数が知りたい場合は $\mathbf{R}=\mathbb{Z}[y]/(y^D-1)$ とすればよい。
TDPC Eのサンプルは以下のように計算できる(0を数えているので1多い)。

D = 7
N = 123456789012345678901234567890
mod = 10^9 + 7
RRR.<y> = GF(mod)[]
RR = QuotientRing(RRR, y^D -1)
R.<x> = RR[]

print(Nth_digital(N, R(1), (1+x+x^2+x^3+x^4+x^5+x^6+x^7+x^8+x^9)*(1+y*x+(y*x)^2+(y*x)^3+(y*x)^4+(y*x)^5+(y*x)^6+(y*x)^7+(y*x)^8+(y*x)^9), 10)[0])
>>> 468357805

算術計算量は $\mathsf{M}(1,18)\lceil\log_{10} (N+1)\rceil$ である。しかし $\mathbf{R}$ の乗算に $O(D^2)$ 時間かけると遅すぎるので、実装には工夫が必要である。多項式 $Q_1(x):=\sum_{n=0}^9 (yx)^n$ の係数が単項式であることを利用して、$P$ と $Q_1$ の乗算に含まれる $\mathbf{R}$ 上の乗算を $O(D)$ 時間でおこない、その後に $\sum_{n=0}^9 x^n$ を掛けるとよい。後に説明するオートマトンを載せたデジタル冪級数の考え方と同じである。実装例

その他の例

基本的には「$N$ 以下の非負整数の中で制約を満たすものの個数を計算せよ」という問題である。制約の種類によっていくつかのタイプに分類している。

桁和タイプ

10進数表現した時の桁の和や登場回数について制約がある。

ABC154 E, Almost Everywhere Zero

概要
$N$以下の非負整数で10進数表現に非零の桁を$K$個含むものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}[y]/y^{K+1}$

Q(x) = \left(\sum_{n=0}^9 x^n\right)\left(1 + y\sum_{n=1}^9 x^n\right).

解は $[y^K]\,[x^N] \prod_{k\ge0} Q(x^{10^k})$である。

ABC029 D, 1

概要
$N$以下の非負整数の10進数表現に現れる1の個数の総和を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}[y]$

Q(x) = \left(\sum_{n=0}^9 x^n\right)\left(1 + yx + \sum_{n=2}^9 x^n\right).

$g(y) := [x^N] \prod_{k\ge0}Q(x^{10^k})$ とすると、解は $g'(1)$ である。

Codeforces1036 C, Classy Numbers

概要
$N$以下の非負整数で10進数表現に非零を高々3個含むものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}[y]/y^4$

Q(x) = \left(\sum_{n=0}^9 x^n\right)\left(1 + y \sum_{n=1}^9 x^n\right).

$g(y) := [x^N] \prod_{k\ge0} Q(x^{10^k})$ とすると、解は $g(1)$ である。

ABC208 E, Digit Products

概要
$N$以下の非負整数で10進数表現の桁積が $K$ 以下であるものの個数を計算せよ。

解法

$\mathbf{R}$ を $\mathbb{Z}$ 上の形式的ディリクレ級数環とする。

\begin{align}
Q(x) &= \left(\sum_{n=0}^9 x^n\right)\left(\sum_{n=0}^9 n^{-s} x^n\right)\\
Q_1(x) &= \left(\sum_{n=0}^9 x^n\right)\left(\sum_{n=1}^9 n^{-s} x^n\right)
\end{align}

$g(s) := [x^N] Q_1(x^{10^{\lfloor\log_{10}N\rfloor}})\prod_{k=0}^{\lfloor\log_{10}N\rfloor-1} Q(x^{10^k})$ とすると $\sum_{t=0}^K [t^{-s}]\, g(s)$ は解のうち桁数が $N$ と等しいものの個数である。
実際には計算の中で登場する形式的ディリクレ級数は非常に疎であり、$t=0$ 及び 2,3,5,7 しか素因数を持たない $t$ についてのみ $t^{-s}$ の係数が非零になる。この問題を解くためにはこの疎性に基づいた高速化が必要である。

SUM-XOR タイプ

いくつかの整数の組の和として表される数で、その整数の組の排他的論理和に制約がある。最終的な冪級数はシンプルであるが、導出までに考察が必要なことが多い。
以下 $\oplus$ は排他的論理和を表す。

ABC129 E, Sum Equals Xor

概要
非負整数の組 $(a,b)$ で $a+b=a\oplus b\le L$ を満たすものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}$
$a+b= (a \oplus b) + 2(a \land b)$ より、

a+b = a \oplus b \iff a \land b = 0

よって、

Q(x) = (1+x)(1+2x)

について、解は $[x^L]\, \prod_{k\ge0} Q(x^{2^k})$ である。

ARC116 D, I Wanna Win The Game

概要
$N$ 個の非負整数の組 $(a_1,\dotsc,a_N) $で $\sum_{i=1}^Na_i = M,\,\bigoplus_{i=1}^N a_i=0$ を満たすものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}$
XOR の条件を無視すると、解は $[x^M]\, \frac1{(1-x)^N} = [x^M]\, \prod_{k\ge0} (1+x^{2^k})^N$ となる。
XOR の条件を考慮すると、$(1+x)^N$ をその偶数次だけからなる多項式に起き替えればよいことが分かる。つまり、

Q(x) = \sum_{n=0}^{\lfloor N/2\rfloor} \binom{N}{2n} x^{2n}

について、解は $[x^M]\, \prod_{k\ge0} Q(x^{2^k})$ である。

ABC050 D, Xor Sum

概要
$N$ 以下の非負整数の組 $(u,v)$ で $a\oplus b=u,\,a+b=v$ を満たすような非負整数 $a,b$が存在するものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}$
整数 $a$ と $b$ について、$a$ の $i$ ビット目が0で $b$ の $i$ ビット目が1のとき、
それらのビットを入れ替えて得られる $a'$ と $b'$ について和とXORは保存されている ($a+b=a'+b'$, $a \oplus b=a'\oplus b'$)。
逆にこれらを重複して数えなければ十分である。このことは

\begin{align}
\Phi\colon\{0,1,2\}^t&\to\{0,1\}^t\times \mathbb{Z}_{\ge0}\\
(x_0,\dotsc,x_{t-1})&\mapsto \left((x_0\bmod2,\dotsc,x_{t-1}\bmod2),\, \sum_{k=0}^{t-1} x_{k} 2^k\right)
\end{align}

が単射であることから分かる。$\Phi$ が単射であることを証明する。$\Phi(x_0,\dotsc,x_{t-1})=\Phi(y_0,\dotsc,y_{t-1})$ とする。このとき $\sum_{k=0}^{t-1} (x_k-y_k)2^k = 0$ を満たす。一方、$x_k=y_k\bmod 2$ より、$x_k-y_k\in\{-2,0,+2\}$ である。ここで、

\begin{align}
A_+ &:= \{k\in\{0,\dotsc,t-1\}\mid x_k - y_k = +2\}\\
A_- &:= \{k\in\{0,\dotsc,t-1\}\mid x_k - y_k = -2\}\\
A_0 &:= \{k\in\{0,\dotsc,t-1\}\mid x_k - y_k = 0\}
\end{align}

と定義すると、$\sum_{k\in A_+} 2^k = \sum_{k\in A_-} 2^k$ が成り立つ。よって、$A_+=A_-=\varnothing$, $A_0 = \{0,\dotsc,t-1\}$ であり、$(x_0,\dotsc,x_{t-1})=(y_0,\dotsc,y_{t-1})$ が得られる。
よって、$(x_k\in\{0,1,2\};{k\ge 0})$ で $\sum_{k\ge0} x_k2^k\le N$ であるものの個数が解となる。つまり、

Q(x) = (1+x)(1+x+x^2)

について、解は $[x^N]\, \prod_{k\ge0} Q(x^{2^k})$ である。
3つの整数について考える場合 ($u=a\oplus b\oplus c$, $v=a+b+c$ ) も同様の考察から $Q(x) = (1+x)(1+x+x^2+x^3)$ とおけばよいことが分かる。

ABC117 D, XXOR

概要
$\max_{0\le a\le K} \sum_{i=1}^N (a \oplus A_i)$ を計算せよ。
このカテゴリに入れるのは強引かもしれない。

解法

$\mathbf{R}=(\max,+)$-半環
$A_s[i]$ を $A_s$ の $i$ ビット目以外を 0 にしたもの, $\overline{A_s[i]}$ を $A_s[i]$ の $i$ ビット目を反転したものとする。

\begin{align}
%\sum_{k=0}^K \prod_{s=1}^N (k \mathop{\mathrm{XOR}} A_s)
%&=
%\sum_{k=0}^K \prod_{s=1}^N \prod_{i=0}^{t-1} [(k_i \mathop{\mathrm{XOR}} A_s[i]) \times{2^{i}}]\\
%&=
[x^K]\,\prod_{i=0}^{t-1} (0+0x^{2^i})\left( \prod_{s=1}^N A_s[i] + x^{2^i}\prod_{s=1}^N \overline{A_s[i]}\right).
\end{align}

実装例

分割タイプ

$m$ 冪の数だけを用いた分割数。

ABC155 E, Payment

概要
$N$ を 10 冪の数の重複無制限の和と差によって表現するとき、項の個数の最小値を計算せよ。

解法

$\mathbf{R}=(\min,+)$-半環

Q(x) = 0+\sum_{n=1}^9 (nx^n+nx^{-n})

について、解は $[x^N]\, \prod_{k\ge0} Q(x^{10^k})$ である。
$P(x) = a + b x^{-1}$ を更新していくことになる。
実装例

ABC231 E, Minimal payments

概要
$X$ と $1=A_1 < A_2 <\cdots < A_n$ が与えられる。ここで、$i=1,\cdots,n-1$ について $A_i \,|\, A_{i+1}$ である。$X$ を $(A_i)_{i=1,\cdots n}$ の重複無制限の和と差によって表現するとき、項の個数の最小値を計算せよ。

解法

$\mathbf{R}=(\min,+)$-半環

各 $i=0,\cdots,n-2$ について $B_i:=A_{i+2}/A_{i+1}$ とおく。そして、各 $i=0,\cdots,n-2$ について

Q_i(x) = 0+\sum_{n=1}^{B_i-1} (nx^n+nx^{-n})

とし、

Q_{n-1}(x) = 0+\sum_{n\ge1} (nx^n+nx^{-n})

とおくと、解は $[x^X]\, \prod_{k=0}^{n-1} Q_i(x^{A_{k+1}})$ である。
実装例

ABC182 F, Valid payments

概要
$X$ と $1=A_1 < A_2 <\cdots < A_n$ が与えられる。ここで、$i=1,\cdots,n-1$ について $A_i \,|\, A_{i+1}$ である。このとき、任意の非負整数 $W$ について、$W$ を$(A_i)_{i=1,\cdots n}$ の重複無制限の和によって表現するとき、最小項数の表現が一意に定まる。$X=Y-Z$ を満たすような非負整数の組 $(Y, Z)$ であって、$Y$ と $Z$ の最小項数表現に共通の整数が現われないものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}$

各 $i=0,\cdots,n-2$ について $B_i:=A_{i+2}/A_{i+1}$ とおく。そして、各 $i=0,\cdots,n-2$ について

Q_i(x) = 1+\sum_{n=1}^{B_i-1} (x^n+x^{-n})

とし、

Q_{n-1}(x) = 1+\sum_{n\ge1} (x^n+x^{-n})

とおくと、解は $[x^X]\, \prod_{k=0}^{n-1} Q(x^{A_{k+1}})$ である。
実装例

TTPC2015 F, レシート

概要
非負整数 $A$ が与えられる。$A=X-Y$ を満たす非負整数 $X$ と $Y$ について $A$ と $X$ と $Y$ の10進数表現の $i$ 桁目が等しくなるような $i$ の個数の最大値を計算せよ。

解法

$\mathbf{R}=(\max,+)$-半環

Q_0(x) = 1+\sum_{n=1}^{9} (0x^n+0x^{-n})

とし、

Q_1(x) = 0+\sum_{n=1}^{9} (0x^n+0x^{-n})

とおくと、解は $[x^A]\, Q_1(x^{\lfloor\log_{10}A\rfloor+1})\prod_{k=0}^{\lfloor\log_{10}A\rfloor} Q_0(x^{10^k})$ である。
実装例

整数タイプ

10進数表現したときに登場する数に制限がある。また、剰余に関する制約がある。

Donuts2014 B, 7th

概要
$N$ 以下の非負整数で10進数表現に 1,2,3,4,5,6,7 のみを含むもののうち 7の倍数であるものの個数を計算せよ。

解法

$\mathbf{R}=\mathbb{Z}$
0 が許されていないので多少面倒である。まず、$N \leftarrow \lfloor N\rfloor_7 :=7\lfloor N/7\rfloor$ と 7の倍数に切り下げておく。

Q(x) = \left(\sum_{n=0}^9 x^{7n}\right)\left(\sum_{n=1}^7 x^n\right) = \sum_{n=1}^{70} x^n

について、 $g(N) := [x^{N}]\, \prod_{k=0}^{\lfloor\log_{10} N\rfloor} Q(x^{10^k})$ は条件を満たすもののうち、桁数が $N$ と同じものだけを数えている。すべての桁数について条件を満たす非負整数を数えるためには、

g(7)+ g(77) + \dotsb + g(\overbrace{7\cdots 7}^{\lfloor \log_{10} N\rfloor}) + g(N)

を計算すればよい。
$g(N)$ は

Nth_digital(N, R(1), Q, 10)

で計算できる。
ここで、$g(7)+ g(77) + \dotsb + g({7\cdots 7})$ の部分は Nth_digital の中のループの各ステップで $P(0)$ の和を取ることで計算することもできる。

整数ペアタイプ

非負整数のペア $(u, v)$ について $u$ と $v$ の $m$ 進数表現の $i$ 桁目 $(u_i, v_i)$ について制約がある。

ABC138 F, Coincidence

概要
非負整数のペア $u, v\in[L, R]$ で $u$ と $v$ の二進数表現の桁数が等しく、$u\land v = v$ を満すものの個数を計算せよ。

解法

以下 noshi91 さん発案の解法である。
$\mathbf{R}=\mathbb{Z}$
桁数の条件を無視すると、

Q(x,y) = (1+x)(1+y)(1+x+xy)

について、 $g(L,R) := [x^{R}y^{L}]\, \prod_{k\ge0} Q(x^{2^k},y^{2^k})$ は $u\land v = v$ を満たし、$u\le R$, $v\le L$ である $(u,v)$ の個数である。
さらに、$h(L,R) := [x^{R-2^{\lfloor\log R\rfloor}}y^{L-2^{\lfloor\log R\rfloor}}]\, \prod_{k=0}^{\lfloor\log R\rfloor-1} Q(x^{2^k},y^{2^k}) $ は $u$ と $v$ の桁数が $R$ の桁数と等しく、 $u\land v = v$ を満たし、$u\le R$, $v\le L$ である $(u,v)$ の個数である。
よって、$w(L,R):=h(R,R) - h(L-1,R)$ は条件を満たす $(u,v)$ の中で、桁数が $R$ と等しいものを数えている。
すべての桁数について条件を満たす非負整数を数えるためには、

w(L,2^{\lceil\log (L+1)\rceil}-1) + w(L,2^{\lceil\log (L+1)\rceil+1}-1) + \dotsb + w(L, 2^{\lfloor\log R\rfloor-1}-1) + w(L,R)

を計算すればよい。実装例

実際には $R$ と $L-1$ の桁数が等しい場合のみ二変数のデジタル冪級数が必要となる。それぞれの固定された桁数について、上界だけが与えられている場合は $(1+x)(1+2x)$ を使って数え上げられるし、下界だけが与えられている場合も $(1+x)(2+x)$ を使って数え上げられる。上界も下界も与えられていない場合は 3 の冪となる。$R$ と $L-1$ の桁数が等しい場合のみ、一つの桁数の中で上下界が同時に与えられ二変数のデジタル冪級数が必要となる。

オートマトンを載せたデジタル冪級数

桁dp というのは整数を $m$ 進数表現した文字列に対する決定性有限オートマトンに基づき、条件を満たす整数の個数を数え上げるアルゴリズムである(オートマトン上の DP (桁 DP の一般化))。任意の桁dp について、対応するデジタル冪級数が存在することを示す。整数に対する条件のうち、「$N$ 以下」と「$D$ の倍数」を除いたものを記述した決定性有限オートマトンについて考える。このとき、オートマトンは整数の $m$ 進数表現を下の桁から順番に読むものとする。オートマトンの状態の集合 $\mathcal{S}$ に重みを与えたもの $w: \mathcal{S}\to\mathbf{R}$ とオートマトン上の状態の遷移 $T: \mathcal{S}\to\mathcal{S}$ の積 $wT: \mathcal{S}\to\mathbf{R}$ を

(w T)(q) := \sum_{q'\in T^{-1}(q)} w(q')\qquad\forall q\in\mathcal{S}

と定義する。そして $\mathcal{S}\to\mathbf{R}$ を係数に持つ多項式 $P$ に対して、$\mathcal{S}\to\mathcal{S}$ を係数に持つ多項式 $Q$ を掛けることで数え上げをおこなう。「$N$ 以下」と「$D$ の倍数」という制約は形式的冪級数を使って代数的に表現される。まず $P$ を定数項だけの多項式とし、その係数をオートマトンの初期状態に重み 1 それ以外に重み 0 とする。そして、

Q(x) = \left(\sum_{i=0}^{m-1} I x^{Di}\right)\left(\sum_{i=0}^{m-1} T_i x^i\right)

とすると $[x^N]\,P(x)Q(x)Q(x^m)Q(x^{m^2})\dotsm$ の受理状態の重みの総和が

1. $N$ 以下で
2. $D$ で割った余りが $N\bmod D$ と等しく
3. $m$ 進数表現がオートマトンで受理される

非負整数の($\mathbf{R}$ における)個数となる。ここで、$I$ は恒等写像、$T_i$ は文字 $i$ に対応するオートマトンの遷移とする。
遷移を係数として持つ多項式同士の乗算も考えられるが、アルゴリズムの中では実装する必要はない。実際に遷移の合成を考えるのも遷移の和を考えるのも面倒である。遷移の行列表現を考えれば自然に合成と和が表現できるが、行列表現は冗長なのであまりよい方法ではない。$P$ に $Q$ を掛けてセクションオペレータを適用するためには先に $P$ に $\sum_{i=0}^{m-1} T_i x^i$ を掛けた後に $\sum_{i=0}^{m-1} I x^{Di}$ を掛けるのとセクションオペレータを適用するのを同時にやるとよい。詳しくは下の実装を参照のこと。

下の桁から読む決定性有限オートマトンの代わりに「下の桁から読む、受理されるときに受理パスが1つであるような非決定性有限オートマトン」でも構わない。そのようなものは(無駄な 0 も含めて)上の桁から読む決定性有限オートマトンの初期状態と受理状態を入れ、遷移の向きを逆にすることによって得ることができる。上の桁から読む方が上位の桁にある無駄な 0 を扱いやすい。

実装の概要

#include <vector>

using u32 = unsigned int;
using u64 = unsigned long long int;
using state = std::vector<u64>; /* オートマトンによる。場合によっては複雑 */
using trans = u32; /* 基本的に 0, 1, ... m-1 で遷移を表す */

state operator*(const state &lhs, trans rhs){
  /* 本質的にここだけ書けばいい */
}

void operator+=(state &a, const state &b){
  /* 状態の重みの += */
}

using poly_state = std::vector<state>;
using poly_trans = std::vector<trans>;

poly_state operator*(const poly_state &a, const poly_trans &b){
  int da = a.size() - 1;
  int db = b.size() - 1;

  poly_state c(da + db + 1);

  for(int i = 0; i <= da; i++){
    for(int j = 0; j <= db; j++){
      c[i+j] += a[i] * b[j];
    }
  }

  return c;
}

/* この関数の返り値の受理状態の重みを足せば
  1. n 以下
  2. d で割った余りが n % d と等しい
  3. オートマトンで受理される
という条件を満たす非負整数の個数が得られる */
state Nth_digital(u64 n, const poly_state &P0, const poly_trans &Q1, u32 d, u32 m){
  poly_state P = P0;
  while(n){
    poly_state U = P * Q1;
    P.clear();
    /* 1+x^d+x^(2d)+...+x^((m-1)d) を掛けてから Section operator */
    U.resize(U.size() + (m - 1) * d);
    for(unsigned i = n % m; i < U.size(); i += m){
      P.push_back(U[i]);
      for(int j = 1; j < m && j * d <= i; j++) P.back() += U[i - j * d];
    }

    n /= m;
  }
  return P[0];
}

ここで state のデフォルトコンストラクタがすべての状態の重みが 0 のものを生成すると仮定している。
$d$ の倍数を数え上げたい場合はあらかじめ $n$ を $d$ の倍数に切り下げておくか、$P_0$ を $x^{n\bmod d}$ 倍しておくとよい。特に $n$ が文字列で与えられる場合は後者の方法が楽である。
U を resize する無駄を省いて

    for(unsigned i = c - '0'; i < U.size() + (m - 1) * d; i += m){
      int j = std::min(i / d, m - 1);
      if(i - j * d < U.size()){
        P.push_back(U[i - j * d]);
        for(j--; j >= 0 && i - j * d < U.size(); j--) P.back() += U[i - j * d];
      }
      else P.push_back(state());
    }

と書くこともできる。

例

説明が大変なので実装例のみ紹介する。

ARC127 A, Leading 1s

概要
$N$ 以下のすべての非負整数について 10 進数表現の上位に連続する1の個数の総和を計算せよ。

解法
上位の0を考慮に入れないオートマトン

using state = std::vector<u64>;
using trans = bool;

state operator*(const state &lhs, trans rhs){
  state ret;
  if(rhs){
    ret.resize(lhs.size()+1);
    ret[0] = 0;
    for(unsigned i = 0; i < lhs.size(); i++) ret[i+1] = lhs[i];
  }
  else {
    ret.push_back(0);
    for(unsigned i = 0; i < lhs.size(); i++) ret[0] += lhs[i];
  }

  return ret;
}

int main(){
  u64 n;
  std::cin >> n;

  poly_state P({{1}});
  poly_trans Q1(10, false);
  Q1[1] = true;

  u64 ans = 0;
  for(u64 t = 10; t <= n; t *= 10){
    state x = Nth_digital(t-1, P, Q1, 1, 10);
    for(unsigned i = 1; i < x.size(); i++) ans += i * x[i];
  }

  state x = Nth_digital(n, P, Q1, 1, 10);
  for(unsigned i = 1; i < x.size(); i++) ans += i * x[i];

  std::cout << ans << std::endl;

  return 0;
}

上位の0も考慮に入れたオートマトン

using state = std::pair<std::vector<u64>, std::vector<u64> >;
using trans = u32;

state operator*(const state &lhs, trans rhs){
  state ret;
  if(rhs == 0){
    ret.second = lhs.first;
    ret.second += lhs.second;
  }
  else if(rhs == 1){
    ret.first.resize(std::max(lhs.first.size()+1, 2UL));
    ret.first[0] = 0;
    for(unsigned i = 0; i < lhs.first.size(); i++) ret.first[i+1] = lhs.first[i];
    for(unsigned i = 0; i < lhs.second.size(); i++) ret.first[1] += lhs.second[i];
  }
  else {
    ret.first.push_back(0);
    for(unsigned i = 0; i < lhs.first.size(); i++) ret.first[0] += lhs.first[i];
    for(unsigned i = 0; i < lhs.second.size(); i++) ret.first[0] += lhs.second[i];
  }

  return ret;
}

int main(){
  u64 n;
  std::cin >> n;

  poly_state P({{{1}, {}}});
  poly_trans Q1(10, 2);
  Q1[0] = 0;
  Q1[1] = 1;

  state x = Nth_digital(n, P, Q1, 1, 10);

  u64 ans = 0;
  for(unsigned i = 1; i < x.first.size(); i++) ans += i * x.first[i];
  for(unsigned i = 1; i < x.second.size(); i++) ans += i * x.second[i];
  std::cout << ans << std::endl;

  return 0;
}

JOI2012 F, ジグザグ数

概要
$N$ 以下の非負整数で 10 進数表現がジグザグしてるものの個数を計算せよ。

実装
汎用性と可読性と速度をある程度両立している。

アルゴリズム(MSB-first)

あまり重要ではない余談にあるように、$N$ の $m$ 進数表現 $N_\ell\dotsm N_0$ について、$L_{N_\ell}\dotsm L_{N_0} P$ の定数項が解となる。$e$ を第一成分が1でそれ以外が0のベクトルとすると、内積 $\langle \cdot,\cdot\rangle$ を用いて

\langle e, L_{N_\ell}\dotsm L_{N_0} P\rangle = 
\langle L_{N_0}^T \dotsm L_{N_\ell}^T e, P\rangle

と解を表すことができる。ここで、$L^T$ は $L$ の転置である。この右辺に基づいて計算するアルゴリズムが MSB-first アルゴリズムである。
特に Bostan–Mori のアルゴリズムの MSB-first バージョンの自然な記述と詳しい説明はX^N mod P(X) の計算で示されている。

二冪分割数

二冪の数を重複無制限で使う分割数の母関数は

\prod_{k\ge0} \frac1{1-x^{2^k}}

である。
このようにデジタル冪級数における多項式 $(Q_k)_{k\ge0}$ が有理多項式であってもデジタル冪級数と見なすことができる。

G(x) := \prod_{k\ge0}\frac{P(x^{2^k})}{Q(x^{2^k})}

とおく。

\begin{align}
[x^N]\, \frac{P_0(x)}{Q_0(x)}G(x) &= [x^N]\, \frac{P_0(x)}{Q_0(x)}\frac{P(x)}{Q(x)}G(x^2)\\
 &= [x^N]\, \frac{P_0(x)P(x)S(-x)}{S(x)S(-x)}G(x^2)\\
 &= \begin{cases}
[x^{N}]\, \frac{U_\mathrm{e}(x^2)}{V(x^2)}G(x^2),&\text{if $N$ is even}\\
[x^{N}]\, \frac{xU_\mathrm{o}(x^2)}{V(x^2)}G(x^2),&\text{otherwise}\\
 \end{cases}\\
 &= \begin{cases}
[x^{N/2}]\, \frac{U_\mathrm{e}(x)}{V(x)}G(x),&\text{if $N$ is even}\\
[x^{(N-1)/2}]\, \frac{U_\mathrm{o}(x)}{V(x)}G(x),&\text{otherwise}\\
 \end{cases}
\end{align}

ここで $S(x) := Q_0(x)Q(x)$, $V(x^2):=S(x)S(-x)$, $P_0(x)P(x)S(-x) = U_\mathrm{e}(x^2) + x U_\mathrm{o}(x^2)$ である。
アルゴリズムは以下のようになる。

R.<x> = ZZ[]

def section2(P, t):
  return R(P.list()[t::2])

def Nth_digital_r(N, P, Q, P0, Q0):
  while N > 0:
    S = Q0 * Q
    P0 = section2(P0 * P * S(-x), N % 2)
    Q0 = section2(S * S(-x), 0)
    N //= 2
  return P0(0) / Q0(0)

print(Nth_digital_r(2^50, R(1), 1-x, R(1), R(1)))

特に $N$ の 二冪分割数を計算するアルゴリズムは以下のように書ける。

R.<x> = ZZ[]

def section2(P, t):
  return R(P.list()[t::2])

def bp(N):
  P = R(1)
  Q = R(1)
  while N > 0:
    Q *= 1 + x
    P = section2(P * Q, N % 2)
    N //= 2
  return P(0)

print(bp(2^50))

このアルゴリズムの算術計算量は

\sum_{i=0}^{\lceil\log(N+1)\rceil-1} i + \mathsf{M}(i+1) = O(\mathsf{M}(\log N)\log N)

である。