04.10_数位dp

数位 DP

问题描述

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

数字计数

算法思想

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

代码实现

c++
java
python

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)

区间数字统计

算法思想

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

代码实现

c++
java
python

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)