Hard

题目描述

给你一个长度为 n 的整数数组 digitSum

如果一个长度为 n 的数组 arr 满足以下条件,则认为它是有效的:

  • 0 <= arr[i] <= 5000
  • 数组是非递减的
  • arr[i] 的各位数字之和等于 digitSum[i]

返回不同有效数组的数量。由于答案可能很大,请返回模 10^9 + 7 的结果。

如果数组中每个元素都大于等于前一个元素(如果存在),则称该数组是非递减的。

示例 1:

输入:digitSum = [25,1]
输出:6
解释:
各位数字之和为 25 的数字有 799、889、898、979、988 和 997。
在保持数组非递减的情况下,唯一能在这些值之后出现且各位数字之和为 1 的数字是 1000。
因此,有效数组为 [799, 1000]、[889, 1000]、[898, 1000]、[979, 1000]、[988, 1000] 和 [997, 1000]。
答案是 6。

示例 2:

输入:digitSum = [1]
输出:4
解释:
有效数组为 [1]、[10]、[100] 和 [1000]。
答案是 4。

示例 3:

输入:digitSum = [2,49,23]
输出:0
解释:
在范围 [0, 5000] 内没有各位数字之和为 49 的整数。因此答案是 0。

约束条件:

  • 1 <= digitSum.length <= 1000
  • 0 <= digitSum[i] <= 50

解题思路

这是一道经典的动态规划题目。核心思路是:

预处理阶段:首先遍历所有可能的数字(0到5000),计算每个数字的各位数字之和,并将具有相同数字和的数字分组存储在有序数组中。

动态规划

  • 状态定义:dp[i][j] 表示到第 i 个位置,选择第 j 个候选数字时的方案数
  • 转移方程:对于第 i 个位置的每个候选数字,我们需要统计前一个位置中所有小于等于当前数字的方案数

优化策略

  1. 使用前缀和优化:对前一层的 dp 值构建前缀和数组
  2. 使用二分查找:通过 upper_bound 快速找到满足条件的范围

具体步骤

  1. 预处理:将 0-5000 的所有数字按数字和分组并排序
  2. 初始化:第一个位置的所有候选数字方案数为 1
  3. 状态转移:对每个位置,计算从前一个位置转移过来的方案数
  4. 结果统计:最后一个位置所有方案数的和

时间复杂度主要由预处理和 DP 转移组成,空间复杂度为候选数字的存储空间。

代码实现

class Solution {
public:
    int countArrays(vector<int>& digitSum) {
        const int MOD = 1e9 + 7;
        
        // 预处理:按数字和分组
        vector<vector<int>> groups(51);
        for (int i = 0; i <= 5000; i++) {
            int sum = 0, temp = i;
            while (temp > 0) {
                sum += temp % 10;
                temp /= 10;
            }
            groups[sum].push_back(i);
        }
        
        int n = digitSum.size();
        vector<long long> dp(groups[digitSum[0]].size(), 1);
        
        for (int i = 1; i < n; i++) {
            if (groups[digitSum[i]].empty()) return 0;
            
            vector<long long> newDp(groups[digitSum[i]].size(), 0);
            
            // 构建前缀和
            vector<long long> prefixSum(dp.size() + 1, 0);
            for (int j = 0; j < dp.size(); j++) {
                prefixSum[j + 1] = (prefixSum[j] + dp[j]) % MOD;
            }
            
            for (int j = 0; j < groups[digitSum[i]].size(); j++) {
                int curr = groups[digitSum[i]][j];
                // 找到前一层中 <= curr 的最大索引
                int idx = upper_bound(groups[digitSum[i-1]].begin(), 
                                    groups[digitSum[i-1]].end(), curr) 
                         - groups[digitSum[i-1]].begin();
                newDp[j] = prefixSum[idx];
            }
            dp = newDp;
        }
        
        long long result = 0;
        for (long long x : dp) {
            result = (result + x) % MOD;
        }
        return result;
    }
};
class Solution:
    def countArrays(self, digitSum: list[int]) -> int:
        MOD = 10**9 + 7
        
        # 预处理:按数字和分组
        groups = [[] for _ in range(51)]
        for i in range(5001):
            digit_sum = sum(int(d) for d in str(i))
            groups[digit_sum].append(i)
        
        n = len(digitSum)
        dp = [1] * len(groups[digitSum[0]])
        
        for i in range(1, n):
            if not groups[digitSum[i]]:
                return 0
            
            new_dp = [0] * len(groups[digitSum[i]])
            
            # 构建前缀和
            prefix_sum = [0] * (len(dp) + 1)
            for j in range(len(dp)):
                prefix_sum[j + 1] = (prefix_sum[j] + dp[j]) % MOD
            
            for j, curr in enumerate(groups[digitSum[i]]):
                # 二分查找
                left, right = 0, len(groups[digitSum[i-1]])
                while left < right:
                    mid = (left + right) // 2
                    if groups[digitSum[i-1]][mid] <= curr:
                        left = mid + 1
                    else:
                        right = mid
                new_dp[j] = prefix_sum[left]
            
            dp = new_dp
        
        return sum(dp) % MOD
public class Solution {
    public int CountArrays(int[] digitSum) {
        const int MOD = 1000000007;
        
        // 预处理:按数字和分组
        List<int>[] groups = new List<int>[51];
        for (int i = 0; i <= 50; i++) {
            groups[i] = new List<int>();
        }
        
        for (int i = 0; i <= 5000; i++) {
            int sum = 0, temp = i;
            while (temp > 0) {
                sum += temp % 10;
                temp /= 10;
            }
            groups[sum].Add(i);
        }
        
        int n = digitSum.Length;
        long[] dp = new long[groups[digitSum[0]].Count];
        Array.Fill(dp, 1L);
        
        for (int i = 1; i < n; i++) {
            if (groups[digitSum[i]].Count == 0) return 0;
            
            long[] newDp = new long[groups[digitSum[i]].Count];
            
            // 构建前缀和
            long[] prefixSum = new long[dp.Length + 1];
            for (int j = 0; j < dp.Length; j++) {
                prefixSum[j + 1] = (prefixSum[j] + dp[j]) % MOD;
            }
            
            for (int j = 0; j < groups[digitSum[i]].Count; j++) {
                int curr = groups[digitSum[i]][j];
                int idx = BinarySearch(groups[digitSum[i-1]], curr);
                newDp[j] = prefixSum[idx];
            }
            dp = newDp;
        }
        
        long result = 0;
        foreach (long x in dp) {
            result = (result + x) % MOD;
        }
        return (int)result;
    }
    
    private int BinarySearch(List<int> list, int target) {
        int left = 0, right = list.Count;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (list[mid] <= target) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left;
    }
}
var countArrays = function(digitSum) {
    const MOD = 1000000007;
    const n = digitSum.length;
    
    // Precompute all numbers with their digit sums
    const numbersByDigitSum = new Array(51).fill(null).map(() => []);
    
    for (let num = 0; num <= 5000; num++) {
        let sum = 0;
        let temp = num;
        while (temp > 0) {
            sum += temp % 10;
            temp = Math.floor(temp / 10);
        }
        if (sum <= 50) {
            numbersByDigitSum[sum].push(num);
        }
    }
    
    // dp[i][j] = number of ways to fill first i positions where last number is numbersByDigitSum[digitSum[i-1]][j]
    const memo = new Map();
    
    function dp(pos, lastNum) {
        if (pos === n) return 1;
        
        const key = `${pos},${lastNum}`;
        if (memo.has(key)) return memo.get(key);
        
        let result = 0;
        const candidates = numbersByDigitSum[digitSum[pos]];
        
        for (let num of candidates) {
            if (num >= lastNum) {
                result = (result + dp(pos + 1, num)) % MOD;
            }
        }
        
        memo.set(key, result);
        return result;
    }
    
    return dp(0, 0);
};

复杂度分析

复杂度类型分析
时间复杂度O(5001 + n × m × log m),其中 n 是数组长度,m 是候选数字的平均数量。预处理需要 O(5001),每个位置的转移需要 O(m × log m)
空间复杂度O(51 × m),其中 m 是每组候选数字的平均数量,用于存储分组信息和 DP 状态