文章

算法刷题系列(五十三):数位 DP——按数字位统计受上界约束的方案

算法刷题系列(五十三):数位 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)。

七、边界与常见错误

  1. 把前导零当作真实位:会对位数相关约束或数字统计造成误计。
  2. 忘记 tight 约束:自由选择所有位会统计到 10^L - 1,而不只是 N。
  3. tight 更新比较错对象:要与当前位的上界 limit 比较,不是与全局 N 比较。
  4. 零本身是否计入需要说清:本文统计 [0,N],数字 0 有效但没有数字 1;若统计正整数可从区间结果中排除 0。
  5. 复用不同上界的缓存:缓存依赖 digits,更换上界必须清空缓存或将其放在每次调用自己的闭包中。
  6. 把计数类型设得过小:位数增长时方案数会迅速增加,Java 使用 long,更大范围需要 BigInteger。

小结

数位 DP 把大范围枚举转成逐位构造。tight 保证不超过上界,started 处理前导零;返回方案数量与累计贡献,可以处理“统计某数字出现次数”这类加总问题。

本文由作者按照 CC BY 4.0 进行授权