牛客题解 | 04.10_数位dp

数位 DP

问题描述

数位动态规划是指统计满足某些条件的数字个数的问题。
常见的有统计区间内数字出现次数、数字计数等。
通常需要按位考虑，从高位到低位进行状态转移。

数字计数

算法思想

1. 定义dp[pos][sum][limit]表示当前第pos位，已有和为sum，是否有限制的方案数
2. 对于每一位，枚举可以填的数字
3. 状态转移时需要考虑是否受到上限的限制
4. 最终将所有合法方案数相加

代码实现

class Solution {
    vector<int> digits;
    vector<vector<vector<int>>> dp;
    
    int dfs(int pos, int sum, bool limit) {
        if (pos == -1) return sum == target;
        if (!limit && dp[pos][sum][0] != -1) return dp[pos][sum][0];
        
        int up = limit ? digits[pos] : 9;
        int ans = 0;
        for (int i = 0; i <= up; i++) {
            ans += dfs(pos - 1, sum + i, limit && i == up);
        }
        
        if (!limit) dp[pos][sum][0] = ans;
        return ans;
    }
    
public:
    int countDigitSum(int n) {
        digits.clear();
        while (n) {
            digits.push_back(n % 10);
            n /= 10;
        }
        
        int m = digits.size();
        dp.assign(m, vector<vector<int>>(1000, vector<int>(2, -1)));
        return dfs(m - 1, 0, true);
    }
};

class Solution {
    int[] digits;
    int[][][] dp;
    
    int dfs(int pos, int sum, boolean limit) {
        if (pos == -1) return sum == target ? 1 : 0;
        if (!limit && dp[pos][sum][0] != -1) return dp[pos][sum][0];
        
        int up = limit ? digits[pos] : 9;
        int ans = 0;
        for (int i = 0; i <= up; i++) {
            ans += dfs(pos - 1, sum + i, limit && i == up);
        }
        
        if (!limit) dp[pos][sum][0] = ans;
        return ans;
    }
    
    public int countDigitSum(int n) {
        digits = new int[32];
        int len = 0;
        while (n > 0) {
            digits[len++] = n % 10;
            n /= 10;
        }
        
        dp = new int[len][1000][2];
        for (int i = 0; i < len; i++)
            for (int j = 0; j < 1000; j++)
                Arrays.fill(dp[i][j], -1);
                
        return dfs(len - 1, 0, true);
    }
}

def countDigitSum(n: int) -> int:
    digits = []
    while n:
        digits.append(n % 10)
        n //= 10
    
    @lru_cache(None)
    def dfs(pos: int, sum: int, limit: bool) -> int:
        if pos == -1:
            return 1 if sum == target else 0
            
        up = digits[pos] if limit else 9
        ans = 0
        for i in range(up + 1):
            ans += dfs(pos - 1, sum + i, limit and i == up)
        
        return ans
    
    return dfs(len(digits) - 1, 0, True)

区间数字统计

算法思想

1. 定义dp[pos][cnt][limit]表示当前第pos位，数字出现cnt次，是否有限制的方案数
2. 对于每一位，枚举可以填的数字，更新出现次数
3. 状态转移时需要考虑前导零和上限的限制
4. 最终将所有合法方案数相加

代码实现

class Solution {
    vector<int> digits;
    vector<vector<vector<int>>> dp;
    int digit;
    
    int dfs(int pos, int cnt, bool limit) {
        if (pos == -1) return cnt;
        if (!limit && dp[pos][cnt][0] != -1) return dp[pos][cnt][0];
        
        int up = limit ? digits[pos] : 9;
        int ans = 0;
        for (int i = 0; i <= up; i++) {
            ans += dfs(pos - 1, cnt + (i == digit), limit && i == up);
        }
        
        if (!limit) dp[pos][cnt][0] = ans;
        return ans;
    }
    
public:
    int countDigit(int n, int d) {
        digits.clear();
        while (n) {
            digits.push_back(n % 10);
            n /= 10;
        }
        
        digit = d;
        int m = digits.size();
        dp.assign(m, vector<vector<int>>(m + 1, vector<int>(2, -1)));
        return dfs(m - 1, 0, true);
    }
};

class Solution {
    int[] digits;
    int[][][] dp;
    int digit;
    
    int dfs(int pos, int cnt, boolean limit) {
        if (pos == -1) return cnt;
        if (!limit && dp[pos][cnt][0] != -1) return dp[pos][cnt][0];
        
        int up = limit ? digits[pos] : 9;
        int ans = 0;
        for (int i = 0; i <= up; i++) {
            ans += dfs(pos - 1, cnt + (i == digit ? 1 : 0), limit && i == up);
        }
        
        if (!limit) dp[pos][cnt][0] = ans;
        return ans;
    }
    
    public int countDigit(int n, int d) {
        digits = new int[32];
        int len = 0;
        while (n > 0) {
            digits[len++] = n % 10;
            n /= 10;
        }
        
        digit = d;
        dp = new int[len][len + 1][2];
        for (int i = 0; i < len; i++)
            for (int j = 0; j <= len; j++)
                Arrays.fill(dp[i][j], -1);
                
        return dfs(len - 1, 0, true);
    }
}

def countDigit(n: int, d: int) -> int:
    digits = []
    while n:
        digits.append(n % 10)
        n //= 10
    
    @lru_cache(None)
    def dfs(pos: int, cnt: int, limit: bool) -> int:
        if pos == -1:
            return cnt
            
        up = digits[pos] if limit else 9
        ans = 0
        for i in range(up + 1):
            ans += dfs(pos - 1, cnt + (1 if i == d else 0), limit and i == up)
        
        return ans
    
    return dfs(len(digits) - 1, 0, True)

时间复杂度分析

算法	时间复杂度	空间复杂度
数字计数	$\mathcal{O}(n\cdot m\cdot k)$	$\mathcal{O}(n\cdot m)$
区间统计	$\mathcal{O}(n\cdot m\cdot k)$	$\mathcal{O}(n\cdot m)$

其中n是数位个数，m是状态数，k是每位可以填的数字个数。

应用场景

1. 区间数字统计
2. 数字计数问题
3. 数位和问题
4. 数字特征统计
5. 区间内特殊数

注意事项

1. 前导零的处理
2. 限制条件的判断
3. 状态定义的选择
4. 记忆化的实现
5. 边界情况的处理

经典例题

1. 数字1的个数
2. 数字游戏
3. 数字计数
4. 串