Digit DP: как посчитать числа на отрезке [L, R], у которых сумма цифр делится на k?
Усложнение цифрового ДП: на отрезке $[L, R]$ ($R \le 10^{18}$) нужно число тех, у кого сумма цифр кратна $k$ ($k \le 100$). Как добавить «сумму цифр по модулю» в состояние и как корректно обработать диапазон, а не только [0, N]?
2 ответа
Добавляем в состояние остаток суммы цифр по модулю k: go(pos, rem, tight), где rem = (сумма уже поставленных цифр) mod k. В конце засчитываем число, если rem == 0.
Диапазон считаем как f(R) - f(L - 1), где f(X) = количество подходящих чисел в [0, X].
string s;
int K;
long long memo[20][105]; // только для tight=false
bool vis[20][105];
long long go(int pos, int rem, bool tight) {
if (pos == (int)s.size()) return (rem == 0) ? 1 : 0;
if (!tight && vis[pos][rem]) return memo[pos][rem];
int hi = tight ? (s[pos] - '0') : 9;
long long res = 0;
for (int d = 0; d <= hi; ++d)
res += go(pos + 1, (rem + d) % K, tight && (d == hi));
if (!tight) { vis[pos][rem] = true; memo[pos][rem] = res; }
return res;
}
long long f(long long X) {
if (X < 0) return 0;
s = to_string(X);
memset(vis, 0, sizeof vis);
return go(0, 0, true);
}
// ответ: f(R) - f(L - 1)
Состояний $O(\text{len} \cdot k)$, переход $O(10)$ → время $O(\text{len}\cdot k\cdot 10)$, память $O(\text{len}\cdot k)$. Для len=19, k=100 — копейки.
Грабли с memset(vis,...) внутри f: его нужно делать на каждый новый X, потому что значения memo зависят от длины записи X (она задаёт s.size(), до которой мы доходим). Если посчитать f(R) и f(L-1) без сброса — словишь WA, т.к. кеш от R некорректен для другой длины. Безопаснее всего сбрасывать vis в начале каждого f.
Если в состоянии много полей (pos, rem, tight, leading_zero, ...), часто удобнее не глобальные массивы, а map от кортежа или явный многомерный массив — но следи за тем, чтобы кешировать только tight=false.