解説
整数 x の十進表記における桁和を ds(x) とします。
求める値を
S=1≤x≤Nds(x) が素数∑xK
とすると、この問題では SmodPQ を求めます。
P,Q は P<Q を満たす異なる素数なので、
gcd(P,Q)=1
です。
そこで、まず
SP=SmodP
と
SQ=SmodQ
をそれぞれ桁 DP で求めます。
最後に、中国剰余定理を用いて SmodPQ を復元します。
以下では、法を R として、
R=P
または
R=Q
の場合の答えを求める方法を説明します。同じ処理を P,Q についてそれぞれ行います。
1. 補助変数を導入する
整数 x の桁和を
s=ds(x)
とします。
十進整数とその桁和は法 9 で合同なので、
x≡s(mod9)
が成り立ちます。
したがって、
u=9x−s
は整数です。
ここで、現在の整数 x の末尾に数字 d を追加したとします。
新しい整数 x′ と桁和 s′ は、
x′=10x+d,s′=s+d
です。
このとき、新しい補助変数 u′ は、
u′=9x′−s′=910x+d−(s+d)=910x−s=109x−s+s=10u+s
となります。
よって、
u′=10u+s
です。
この式に、追加する数字 d が含まれていないことが重要です。
なお、ここでの (x−s)/9 は整数上の割り算です。法 R における 9 の逆元を使っているわけではありません。
2. 桁 DP の状態
N の桁数を L とします。
L 桁未満の整数についても、先頭に 0 を補って L 桁の整数として扱います。先頭に 0 を補っても桁和は変化しません。
N の先頭から何桁かを処理した時点で、
dpj[s]
を次のように定義します。
- 現在の接頭辞が、N の同じ長さの接頭辞より小さい
- 現在の桁和が s
- 補助変数が u=(x−s)/9
である接頭辞についての uj の総和を表すものとします。
すなわち、
dpj[s]=∑uj(modR)
です。
必要なのは、
0≤j≤K
の範囲だけです。
特に j=0 のときは、各接頭辞が 1 を寄与するものと考えます。そのため、
dp0[s]
は、条件を満たす接頭辞の個数です。
一方、まだ N と完全に一致している接頭辞は常に 1 個しかありません。
そこで、一致している接頭辞については、
equal_sum: 現在の桁和
equal_u: 現在の umodR
の 2 つだけを別に管理します。
3. すでに N より小さい状態からの遷移
現在の桁和が s である接頭辞を考えます。
次の桁を追加した後の補助変数は、
u′=10u+s
です。
したがって、その m 乗は二項定理より、
(u′)m=(10u+s)m=j=0∑m(jm)(10u)jsm−j=j=0∑m(jm)10jsm−juj
となります。
古い桁和が s であるすべての接頭辞について、遷移後の (u′)m を合計した値を Am[s] とすると、
Am[s]=j=0∑m(jm)10jsm−jdpj[s]
です。
この値は、次に追加する数字には依存しません。
次に追加する数字 d は、
d=0,1,…,9
のどれでも選べます。
古い桁和が s なら、新しい桁和は s+d です。
したがって、新しい桁和が t である状態について、
ndpm[t]=0≤d≤9s+d=t∑Am[s]=s=t−9∑tAm[s]
となります。ただし、存在しない桁和は無視します。
よって、
ndpm[t]=Am[t]+Am[t−1]+⋯+Am[t−9]
です。
これは長さ 10 の区間和です。
t を 1 増やすとき、
- 新しく Am[t] を加える
- 範囲外になる Am[t−10] を引く
ことで更新できるため、スライド窓を用いて各 m について線形時間で計算できます。
4. N と一致している状態からの遷移
次に処理する N の数字を a とします。
現在まで N と一致している接頭辞の桁和を seq、補助変数を ueq とします。
次の桁を追加した後の補助変数は、
unext=10ueq+seq
です。
次の数字として d を選ぶとき、
- d<a なら、この桁で初めて N より小さくなる
- d=a なら、引き続き N と一致する
- d>a なら、N より大きくなるため選べない
となります。
d<a を選んだ場合、新しい桁和は、
seq+d
です。
また、unext は d に依存しません。
したがって、各
d=0,1,…,a−1
について、
ndpm[seq+d]+=unextm
とします。
その後、d=a を選んだ一致状態を、
ueq←unext
seq←seq+a
として次の桁へ引き継ぎます。
5. xK を復元する
補助変数の定義より、
x=9u+s
です。
したがって、二項定理から、
xK=(9u+s)K=j=0∑K(jK)(9u)jsK−j=j=0∑K(jK)9jsK−juj
となります。
よって、桁和が s である x<N 全体についての K 乗和は、
j=0∑K(jK)9jsK−jdpj[s]
です。
これを、素数である桁和 s について合計します。
6. N 自身を加える
DP に格納されているのは、
0≤x<N
を満たす整数です。
そのため、最後に N 自身の桁和が素数なら、
NKmodR
を答えに加えます。
一致状態として管理していた equal_u と equal_sum を使うと、
N=9u+ds(N)
なので、
N≡9⋅equal_u+equal_sum(modR)
として NmodR を復元できます。
なお、DP には x=0 も含まれます。しかし、0 の桁和は 0 であり素数ではないため、答えには影響しません。
以上の桁 DP により、法 R における答えを求められます。
この処理を R=P と R=Q についてそれぞれ行い、
SP=SmodP
および
SQ=SmodQ
を求めます。
7. 中国剰余定理による復元
次の連立合同式を満たす S0 を求めます。
S0≡SP(modP)
S0≡SQ(modQ)
まず、
S0=SP+Pt
と置きます。
この式は必ず、
S0≡SP(modP)
を満たします。
もう一方の条件から、
SP+Pt≡SQ(modQ)
すなわち、
Pt≡SQ−SP(modQ)
です。
P,Q は異なる素数なので、P は法 Q で逆元を持ちます。
したがって、
t≡(SQ−SP)P−1(modQ)
です。
0≤t<Q を満たす代表元を選べば、
S0=SP+P((SQ−SP)P−1modQ)
となります。
Q は素数なので、フェルマーの小定理から、
P−1≡PQ−2(modQ)
です。
また、
0≤SP<P,0≤t<Q
なので、
0≤S0<PQ
です。
したがって、この S0 がそのまま SmodPQ の答えになります。
正当性の証明
補題 1
整数 x の桁和を s、補助変数を
u=9x−s
とする。
x の末尾に数字 d を追加した後の補助変数は、
u′=10u+s
である。
証明
新しい整数と桁和は、
x′=10x+d,s′=s+d
なので、
u′=9x′−s′=910x+d−(s+d)=910x−s=109x−s+s=10u+s
です。
よって成立します。□
補題 2
すでに N より小さい接頭辞からの遷移後、ndpm[t] は、新しい桁和が t である接頭辞についての (u′)m の総和になる。
証明
古い桁和が s である状態では、補題 1 より、
u′=10u+s
です。
二項定理から、
(u′)m=j=0∑m(jm)10jsm−juj
なので、古い桁和が s であるすべての状態について合計すると、
Am[s]=j=0∑m(jm)10jsm−jdpj[s]
が得られます。
次の数字 d は 0 から 9 まで自由に選べ、新しい桁和は s+d です。
したがって、新しい桁和が t になるすべての遷移の総和は、
s=t−9∑tAm[s]
です。
これは ndpm[t] の計算と一致します。□
補題 3
各桁の処理後、DP は N の同じ長さの接頭辞より小さい接頭辞を、重複も不足もなく表している。
証明
処理した桁数についての帰納法で示します。
0 桁を処理した時点では、N より小さい接頭辞は存在せず、空の接頭辞が N と一致している状態だけが存在します。
ある桁数まで主張が成立していると仮定します。
すでに N より小さい接頭辞については、次の数字として 0,1,…,9 のすべてを選べます。これらは補題 2 の遷移によってすべて DP に追加されます。
一方、N と一致している接頭辞について、次の N の数字を a とすると、
- d<a を選んだものは、新たに N より小さくなる
- d=a を選んだものは、引き続き N と一致する
- d>a を選んだものは、N を超えるため除外される
となります。
各接頭辞がどの遷移から生成されるかは一意なので、重複も不足もありません。
よって、次の桁を処理した後についても主張が成立します。□
補題 4
calc_mod は、求める総和を法 R で正しく計算する。
証明
補題 3 より、すべての桁を処理した後の DP は、
0≤x<N
を満たすすべての整数を正しく表しています。
各整数について、
x=9u+s
であるため、
xK=j=0∑K(jK)9jsK−juj
と展開できます。
DP は各 uj の総和を保持しているので、この式から、桁和が s である整数の K 乗和を正しく復元できます。
素数である桁和 s だけを合計するため、条件を満たす x<N のみが答えに加算されます。
さらに、N 自身の桁和が素数である場合に限り NK を加えるため、条件を満たすすべての 1≤x≤N がちょうど一度ずつ数えられます。
よって、calc_mod は求める総和を法 R で正しく計算します。□
補題 5
中国剰余定理によって復元した値 S0 は、求める総和と法 PQ で合同である。
証明
構成から、
S0≡SP≡S(modP)
および、
S0≡SQ≡S(modQ)
です。
P,Q は互いに素なので、中国剰余定理より、この連立合同式の解は法 PQ で一意です。
したがって、
S0≡S(modPQ)
です。□
定理
このアルゴリズムは、桁和が素数である N 以下の正整数の K 乗和を、法 PQ で正しく求める。
証明
補題 4 より、桁 DP によって SP と SQ を正しく求められます。
さらに補題 5 より、それらを中国剰余定理で復元した値は、求める総和と法 PQ で合同です。
よって、アルゴリズムは正しいです。□
計算量
L=∣N∣ とします。
i 桁を処理した時点で、桁和は高々 9i です。そのため、各桁で扱う桁和状態数は O(L) です。
各桁和 s について、すべての m,j の組を処理する回数は、
m=0∑K(m+1)=O(K2)
です。
したがって、法を 1 つ固定したときの時間計算量は、
O(K2L2)
です。
これを P,Q の 2 回行いますが、定数倍が増えるだけなので、全体の時間計算量も、
O(K2L2)
です。
中国剰余定理で必要な逆元の計算には O(logQ) 時間かかります。
遷移係数を事前計算するため、メモリ使用量は、
O(K2L)
です。
P と Q に対する桁 DP は順番に実行するため、両方の DP 配列を同時に保持する必要はありません。
制約では、
L≤300,K≤10
です。
想定解コード(Python)