二分数字组合|豆包MarsCode AI刷题

232 阅读4分钟

01问题描述

小F面临一个有趣的挑战:给定一个数组,她需要将数组中的数字分为两组。分组的目标是使得一组数字的和的个位数等于给定的 A,另一组数字的和的个位数等于给定的 B。除此之外,还有一种特殊情况允许其中一组为空,但剩余数字和的个位数必须等于 A 或 B。小F需要计算所有可能的划分方式。

例如,对于数组 [1, 1, 1] 和目标 A = 1,B = 2,可行的划分包括三种:每个 1 单独作为一组,其余两个 1 形成另一组。如果 A = 3,B = 5,当所有数字加和的个位数为 3 或 5 时,可以有一组为非空,另一组为空。


测试样例

样例1:

输入:n = 3,A = 1,B = 2,array_a = [1, 1, 1] 输出:3

样例2:

输入:n = 3,A = 3,B = 5,array_a = [1, 1, 1] 输出:1

样例3:

输入:n = 2,A = 1,B = 1,array_a = [1, 1] 输出:2

样例4:

输入:n = 5,A = 3,B = 7,array_a = [2, 3, 5, 7, 9] 输出:0

02思路

动态规划

2.1预处理

先将数组每个数都模10,只关心个位数

for i in range(n):
    array_a[i] %= 10

2.2 计算 total_sum 并检查特殊情况

total_sum = sum(array_a) % 10

2.2.1一组为空

if total_sum == A or total_sum == B:
    return 1

一组加起来等于A或B,另一组为空,一种划分

2.2.2无法划分

if total_sum != (A + B) % 10:
    return 0

2.3动态规划

2.3.1初始化

f[i][j] 表示用前 i 个元素组成的子集和模10为 j 的组合数。初始条件 f[0][0] = 1 表示用0个元素时,和为0的组合数为1。

i从0到n,j可取0到9

f[i][j]为n+1行,10列

f = [[0] * 10 for _ in range(n + 1)]
f[0][0] = 1

2.3.2动态规划转移方程

for i in range(1, n + 1):
    for j in range(10):
        f[i][j] += f[i - 1][j]
        f[i][j] += f[i - 1][(j - array_a[i - 1] + 10) % 10]

计算前i 个元素组成的子集和模10为 j 的组合数有两种可能情况,这两种情况都要考虑,才能包含所有的组合数

  1. 不选第i个元素,保持上一个状态的值不变

f[i][j] += f[i - 1][j]

  1. 选择第i个元素,将它的个位数加入到目标和中

f[i][j] += f[i - 1][(j - array_a[i - 1] + 10) % 10]

新的目标个位数为 j,而上一步的和的个位数则应当为 j - array_a[i - 1],选了第i个,加上个位数就变成j了,为了避免负数结果,计算 (j - array_a[i - 1] + 10) % 10

2.3.3目标状态

return f[n][B]

存储了能使某一组的和的个位数等于 B 的方案数

03代码

def solution(n, A, B, array_a):
    # n表示数组长度
    for i in range(n):
        array_a[i] %= 10
    total_sum = sum(array_a) % 10
    # 一组为空
    if total_sum == A or total_sum == B:
        return 1
    # 无法划分
    if total_sum != (A + B) % 10:
        return 0
    f = [[0] * 10 for _ in range(n + 1)]
    f[0][0] = 1
    for i in range(n + 1):
        for j in range(10):
            f[i][j] += f[i - 1][j]
            f[i][j] += f[i - 1][(j - array_a[i - 1] - j + 10) % 10]
    return f[n][B]

f[i][j] += f[i - 1][(j - array_a[i - 1] ``- j`` + 10) % 10]

这里多写了一个-j

最终代码:

def solution(n, A, B, array_a):
    # n表示数组长度
    for i in range(n):
        array_a[i] %= 10
    total_sum = sum(array_a) % 10
    # 一组为空
    if total_sum == A or total_sum == B:
        return 1
    # 无法划分
    if total_sum != (A + B) % 10:
        return 0
    f = [[0] * 10 for _ in range(n + 1)]
    f[0][0] = 1
    for i in range(n + 1):
        for j in range(10):
            f[i][j] += f[i - 1][j]
            f[i][j] += f[i - 1][(j - array_a[i - 1] + 10) % 10]
    return f[n][B]

if __name__ == "__main__":
    #  You can add more test cases here
    array1 = [1, 1, 1]
    array2 = [1, 1, 1]
    array3 = [1, 1]
 
    print(solution(3, 1, 2, array1) == 3)
    print(solution(3, 3, 5, array2) == 1)
    print(solution(2, 1, 1, array3) == 2)

04总结

还是没有掌握动态规划算法,不能自己设出dp数组,对本题动态转移方程的理解也没有太深入。

05参考资料