H - Digit Sum Divisible 解説 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;
}

投稿日時:
最終更新: