H - Digit Sum Divisible Editorial
by
miscalculation53
桁 DP を下の桁から決めていく方法
公式解説は上の桁から決めていく桁 DP で、本問の場合はおそらくそちらのほうが楽だと思います。下の桁から決めていく考え方でも解くことができるので、解説します。
[1] 下の桁から決めていく桁 DP
- \(f(n, s, r, m) := n\) 以下の非負整数のうち、桁和が \(s\) であり、\(m\) で割ったあまりが \(r\) であるものの個数
と定めます。答えは \(\displaystyle \sum_{s = 1}^{9 \times 14} f(N, s, 0, s)\) です。\(f(n, s, r, m)\) を(メモ化再帰を利用した)DP で求めましょう。
一番下の桁 \(d\) を \(10\) 通り試します。このとき、\(v = 10v' + d \: (0 \leq d \leq 9)\) について、
- \(v\) が \(n\) 以下で、桁和が \(s\) で、\(m\) で割ったあまりが \(r\) である
という条件を、
- \(v'\) が \(n'\) 以下で、桁和が \(s'\) で、\(m'\) で割ったあまりが \(r'\) である
という条件に変換できれば、DP の遷移が
\[f(n, s, r, m) \ \text{+=} \ f(n', s', r', m')\]
と書けます。
[2] \(n', s'\) を求める
\[ \begin{aligned} n' &= \begin{cases} \left\lfloor \frac{n}{10} \right\rfloor & \text{if} \ d \leq (n \bmod 10) \\ \left\lfloor \frac{n}{10} \right\rfloor - 1 & \text{if} \ d > (n \bmod 10) \end{cases} \\ \displaystyle s' &= s - d \end{aligned} \]
となります。
[3] \(r', m'\) を求める
\(r', m'\) を求めることは、
\[10v' \equiv r - d \pmod m \tag{*}\]
を \(v'\) について解く、つまり
\[v' \equiv r' \pmod {m'}\]
の形で表すことです。
これは、extgcd(拡張ユークリッドの互除法)と呼ばれるアルゴリズムで求めることができます。下の実装例では、ACL の internal::inv_gcd
を利用しています。inv_gcd(a, b)
は、
\(g = \gcd(a, b)\)
\(ax \equiv g \pmod b, 0 \leq x < \dfrac{b}{g}\)
を満たす \((g, x)\) のペアを返すものです。
2 つ目の条件に関する注
$ax \bmod b$ のとりうる値は $0, g, 2g, \ldots, \left(\dfrac{b}{g} - 1 \right)g$ の $\dfrac{b}{g}$ 通りで、その値は $\dfrac{b}{g}$ 周期で繰り返されます(つまり、$x \equiv x' \: \left( \bmod \ \dfrac{b}{g} \right)$ ならば $ax \equiv ax' \pmod b$)。これは、$\dfrac{a}{g}x \bmod \dfrac{b}{g}$ について考えるとわかります($\dfrac{a}{g}$ と $\dfrac{b}{g}$ が互いに素であることに注意)。\((g, x) = \mathrm{inv\_gcd}(10, m)\) とすると、
\(r - d\) が \(g\) の倍数でないとき、\((*)\) は解なし
\(r - d\) が \(g\) の倍数であるとき、\((*)\) の解は \(\displaystyle v' \equiv \frac{r - d}{g} \cdot x \: \left( \bmod \ \frac{m}{g} \right)\)
となります。
[4] 計算量
計算途中の \(n\) の値としてありうるのは、\(N\) を \(10\) で割っていったものと、(\(-1\) されたことにより現れる)その近辺しかないので、\(O(\log N)\) 個程度です。よって、\(s\) の上限を \(S \: (= 126)\) とすると、状態数は \(O(S^3 \log N)\) です。\(B = 10\) として \(S = O(B \log N)\) で、遷移が \(B\) 通りあるので、計算量は(メモ化や extgcd の部分を無視すれば)\(O((B \log N)^4)\) です。
実装例
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#include <atcoder/internal_math>
using namespace atcoder;
// f(n, s, r, m) := n 以下で、桁和が s で、m で割って r あまるものの個数
map<ll, ll> mem[130][130][130]; // mem[s][r][m][n] にメモ
ll f(ll n, int s, int r, int m)
{
if (n == 0)
return s == 0 && r == 0;
if (mem[s][r][m].contains(n))
return mem[s][r][m][n];
ll res = 0;
for (int d = 0; d < 10; d++)
{
ll nn = n / 10 - (d > n % 10);
int ns = s - d;
if (nn < 0 || ns < 0)
continue;
auto [g, x] = internal::inv_gcd(10, m);
if ((r - d) % g != 0)
continue;
int nm = m / g;
int nr = (((r - d) / g * x) % nm + nm) % nm;
res += f(nn, ns, nr, nm);
}
mem[s][r][m][n] = res;
return res;
}
int main()
{
ll N;
cin >> N;
ll ans = 0;
for (int s = 1; s <= 126; s++)
{
ans += f(N, s, 0, s);
}
cout << ans << endl;
}
posted:
last update: