算法刷题系列(五十三):数位 DP——按数字位统计受上界约束的方案
算法刷题系列(五十三):数位 DP——按数字位统计受上界约束的方案
当上界很大,逐个枚举 0 到 N 不现实;但如果问题只关心每一位的性质,可以从高位到低位构造数字,并用状态统计方案。数位 DP 的两个重要标记是 tight(前缀是否仍贴着上界)与 started(是否已经出现非前导零)。
本文计算 [0, N] 中十进制表示里数字 1 出现的总次数。比如数字 101 贡献 2,数字 11 贡献 2。
一、状态定义
把上界 N 转成数字数组 digits,从最高位开始处理。状态 dfs(pos, tight, started) 表示:
- 当前待填的位置是
pos; tight为真时,前缀与N的前缀相同,当前位不能超过上界对应数字;started为真时,之前已经填过非零数字,当前数字的真实表示已经开始。
每个状态返回一对结果:
1
(ways, ones)
ways 是可形成的数字个数;ones 是这些数字在已处理位置及后续位置上的数字 1 总数。选当前位 digit 后,子问题提供 child_ways 个后缀方案。若当前位是真实数字中的 1,它为每个后缀各贡献一次:
1
ones += child_ones + (child_ways if this position is a real 1 else 0)
前导零不会被计为数字的一部分,但数值 0 本身仍被统计为一个合法数。
二、tight 如何更新
当前位可选范围为 0 到 limit,其中 limit 是上界该位(tight 时)或 9。下一状态仍 tight,当且仅当当前状态 tight 且所选数字恰好等于 limit。一旦某位选得更小,后面的位可以自由取 0..9。
三、例子:统计 [0, 25]
个位为 1 的数字有 1、11、21,贡献 3 次;十位为 1 的数字有 10 到 19,贡献 10 次;合计 13。
前导零不会把 01 当作十位也有一个 1。每个数只按其通常的无前导零表示统计。
四、Java 实现
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import java.util.Arrays;
public class DigitDp {
private static char[] digits;
private static long[][][][] memo;
private static boolean[][][] seen;
public static long countOnes(long upperBound) {
if (upperBound < 0) {
return 0;
}
digits = Long.toString(upperBound).toCharArray();
int n = digits.length;
memo = new long[n][2][2][2];
seen = new boolean[n][2][2];
return dfs(0, 1, 0)[1];
}
private static long[] dfs(int pos, int tight, int started) {
if (pos == digits.length) {
return new long[]{1, 0};
}
if (seen[pos][tight][started]) {
return memo[pos][tight][started];
}
seen[pos][tight][started] = true;
int limit = tight == 1 ? digits[pos] - '0' : 9;
long ways = 0;
long ones = 0;
for (int digit = 0; digit <= limit; digit++) {
int nextTight = tight == 1 && digit == limit ? 1 : 0;
int nextStarted = started == 1 || digit != 0 ? 1 : 0;
long[] child = dfs(pos + 1, nextTight, nextStarted);
ways += child[0];
ones += child[1];
if (nextStarted == 1 && digit == 1) {
ones += child[0];
}
}
memo[pos][tight][started] = new long[]{ways, ones};
return memo[pos][tight][started];
}
public static void main(String[] args) {
if (countOnes(25) != 13 || countOnes(0) != 0
|| countOnes(9) != 1 || countOnes(100) != 21) {
throw new AssertionError("digit DP test failed");
}
}
}
每个位置的上界数字 limit 在 tight == 1 时与当前位上界相同,因此 digit == limit 正是保持 tight 的条件。started 用于说明数字是否开始;当正在填的数字是 1 时它必然已经开始。
五、Python 实现
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from functools import lru_cache
def count_ones(upper_bound):
if upper_bound < 0:
return 0
digits = [int(ch) for ch in str(upper_bound)]
@lru_cache(maxsize=None)
def dfs(pos, tight, started):
if pos == len(digits):
return 1, 0
limit = digits[pos] if tight else 9
ways = 0
ones = 0
for digit in range(limit + 1):
next_tight = tight and digit == limit
next_started = started or digit != 0
child_ways, child_ones = dfs(
pos + 1, next_tight, next_started
)
ways += child_ways
ones += child_ones
if next_started and digit == 1:
ones += child_ways
return ways, ones
return dfs(0, True, False)[1]
assert count_ones(25) == 13
assert count_ones(0) == 0
assert count_ones(9) == 1
assert count_ones(100) == 21
六、正确性与复杂度
从高位到低位,每个 0..N 中的数字都对应唯一的一条逐位选择路径,且每条不超过上界的路径都构成一个合法数。tight 限制保证没有路径超过 N;started 保证前导零不作为真实数字位参与统计。对子问题返回的 ones 求和,并对当前位的 1 乘以该分支后缀方案数,恰好计入每个数字中每个 1 一次。
状态数为 O(L)(L 是十进制位数,tight 和 started 只有两种取值),每个状态最多尝试 10 个数字,时间为 O(10L)、记忆化空间为 O(L)。
七、边界与常见错误
- 把前导零当作真实位:会对位数相关约束或数字统计造成误计。
- 忘记 tight 约束:自由选择所有位会统计到
10^L - 1,而不只是N。 - tight 更新比较错对象:要与当前位的上界
limit比较,不是与全局N比较。 - 零本身是否计入需要说清:本文统计
[0,N],数字 0 有效但没有数字 1;若统计正整数可从区间结果中排除 0。 - 复用不同上界的缓存:缓存依赖
digits,更换上界必须清空缓存或将其放在每次调用自己的闭包中。 - 把计数类型设得过小:位数增长时方案数会迅速增加,Java 使用
long,更大范围需要BigInteger。
小结
数位 DP 把大范围枚举转成逐位构造。tight 保证不超过上界,started 处理前导零;返回方案数量与累计贡献,可以处理“统计某数字出现次数”这类加总问题。