SDE2 · 90-minute OA Coding

Microsoft OA 数位和限定的非递减数组计数:前缀和 DP

每个位置指定十进制数位和,值域 1..5000、长度至多 5000;用结尾值 DP 和前缀累加做到 O(n×5000)。

题目重点:每个位置指定十进制数位和,值域 1..5000、长度至多 5000;用结尾值 DP 和前缀累加做到 O(n×5000)。

题意与假设

这道练习题保留了面经中值得训练的核心约束,并将需要澄清的条件单独列出。现场拿到变体时,先和面试官确认这些条件。

  • 非递减允许相等,即 result[i−1]≤result[i];每个数组按位置序列区分。
  • 数位和按普通十进制表示计算,不把前导零作为数字一部分。
  • 空 requiredSum 不在题目约束中,训练函数返回一个空数组方案;正式输入 n≥1。
  • 代码默认最大值 5000;测试可传较小 max_value 对照穷举,不改变原题值域。

解题路线

  1. 预计算 value=1..5000 的十进制数位和,可用 digitSum[v]=digitSum[v//10]+v%10。
  2. 令 dp_i[v] 为前 i+1 个位置合法且最后一个值恰为 v 的方案数。首位置若 digitSum[v] 命中 requiredSum[0],dp_0[v]=1,否则 0。
  3. 下一位置值 v 必须不小于前值 u,因此 dp_i[v]=Σ_{u≤v} dp_{i−1}[u],但只有 digitSum[v] 符合当前目标时才保留。
  4. 每一层从 v=1..5000 滚动维护 previous 的前缀和,查询转移变 O(1);层间只留两行 DP,所有累加按 MOD 取模。
  5. 答案为末层所有 v 的 dp 值之和。直接枚举前值 u 会到 O(n×5000²),无法利用非递减的前缀结构。

手动走一遍

requiredSum=[1] 时值域 1..5000 内仅 1、10、100、1000 数位和为 1,所以答案 4。若要求 [1,1],可从这四个值中有放回、按非递减取两项,答案 C(4+2−1,2)=10;DP 的各结尾计数依次为 1、2、3、4。

Python 参考实现

下面的代码按上文假设实现,适合本地运行并追加自己的测试用例。

代码区可横向滚动;使用键盘时,先聚焦代码区再按方向键查看。

MOD = 1_000_000_007


def count_arrays(required_sum: list[int], max_value: int = 5000) -> int:
    """Count nondecreasing positive arrays with prescribed per-value digit sums."""
    if max_value < 1 or any(s < 1 for s in required_sum):
        raise ValueError('invalid bounds')
    if not required_sum:
        return 1
    digit_sum = [0] * (max_value + 1)
    for value in range(1, max_value + 1):
        digit_sum[value] = digit_sum[value // 10] + value % 10
    previous = [0] * (max_value + 1)
    for value in range(1, max_value + 1):
        if digit_sum[value] == required_sum[0]:
            previous[value] = 1
    for target in required_sum[1:]:
        current = [0] * (max_value + 1)
        prefix = 0
        for value in range(1, max_value + 1):
            prefix = (prefix + previous[value]) % MOD
            if digit_sum[value] == target:
                current[value] = prefix
        previous = current
    return sum(previous) % MOD


if __name__ == '__main__':
    assert count_arrays([1]) == 4

复杂度

设 n 为 requiredSum 长、V=5000。数位和预处理 O(V),每层前缀扫描 O(V),总 O(nV) 时间、O(V) 工作空间。n=V=5000 时约 2500 万次位置循环,需避免嵌套枚举所有前值。

易错点与追问

容易失分的边界

  • 把非递减写成严格递增,漏掉重复值方案。
  • 把数位和误当数字本身之和或把 0 加入值域。
  • 朴素转移每个 v 再遍历所有 u≤v,复杂度 O(nV²)。
  • 只对最后答案取模,中间大数在固定宽度语言溢出。
  • 将 的不完整数位和描述当作该题完整样例,编造约束。

继续深挖

  • 如果每个位置有不同上界 Vi,如何限制扫描范围?
  • 若要求严格递增,前缀和取 u<v 还是 u≤v?
  • 若 V 上升到 10^9 但满足特定数位和的候选值稀疏,能否压缩候选集合并用 Fenwick 树?
✓已复制微信号!