青训营X豆包MarsCode 技术训练营第一课 | 豆包MarsCode AI 刷题

76 阅读5分钟

问题描述

小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

解题思路:

问题理解

我们需要将数组中的数字分为两组,使得一组数字的和的个位数等于给定的 A,另一组数字的和的个位数等于给定的 B。还有一种特殊情况允许其中一组为空,但剩余数字和的个位数必须等于 A 或 B。

数据结构选择

由于我们需要处理的是数组的子集和,可以考虑使用动态规划(DP)来解决这个问题。具体来说,我们可以使用一个布尔类型的 DP 数组来记录是否可以得到某个和的个位数。

算法步骤

初始化 DP 数组:创建一个大小为 10 的布尔数组 dp,其中 dp[i] 表示是否可以得到和的个位数为 i。 处理特殊情况:如果数组中所有数字的和的个位数等于 A 或 B,那么有一种划分方式是其中一组为空。 动态规划更新:遍历数组中的每个数字,更新 DP 数组,记录所有可能的和的个位数。 计算划分方式:根据 DP 数组的结果,计算所有可能的划分方式。

`def solution(n, A, B, array_a):
# 计算数组中所有数字的和的个位数
total_sum_mod = sum(array_a) % 10

# 特殊情况处理
if total_sum_mod == A or total_sum_mod == B:
    return 1

# 初始化 DP 数组
dp = [False] * 10
dp[0] = True

# 动态规划更新
for num in array_a:
    # 从后往前更新,避免重复计算
    for i in range(9, -1, -1):
        if dp[i]:
            dp[(i + num) % 10] = True

# 计算划分方式
count = 0
for i in range(10):
    if dp[i]:
        # 检查是否可以形成 A 或 B
        if i == A or i == B:
            count += 1
        # 检查是否可以形成 A 和 B 的组合
        for j in range(10):
            if dp[j] and (i + j) % 10 == A:
                count += 1
            if dp[j] and (i + j) % 10 == B:
                count += 1

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



然后!!就崩了

image.png

由于这个代码是从c++转成python的,我开始去对源代码去修改,修改后的代码段如下:

改进:

规范化数组元素:

for (int& x : array_a) { x %= 10; }

将数组中的每个元素取模10,使得所有元素的值都在0到9之间。

计算总和的模10值:

int total_sum = accumulate(array_a.begin(), array_a.end(), 0) % 10; 使用 accumulate 函数计算数组元素的总和,并取模10。

检查总和是否符合要求:

if (total_sum == A || total_sum == B) { return 1; } if (total_sum != (A + B) % 10) { return 0; } 如果总和的模10值等于 A 或 B,则返回1(意味着可以找到一种组合)。 如果总和的模10值与 A + B 的模10值不相等,则返回0(意味着无法找到组合)。动态规划部分 动态规划数组初始化:

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

动态规划更新:

for (int i = 1; i <= n; ++i) { for (int j = 0; j < 10; ++j) { f[i][j] += f[i - 1][j]; f[i][j] += f[i - 1][(j - array_a[i - 1] + 10) % 10]; } } 外层循环遍历每个元素。 内层循环遍历所有可能的模10值。 f[i][j] 更新包括两种情况:不选第 i 个元素和选择第 i 个元素。 通过对模10值进行取模和偏移,确保数组不越界。 返回结果:

return f[n][B]; 返回用所有 n 个元素组成的子集和模10为 B 的组合数,(确定B的方案那A的方案也就确定了)

最终代码:

`def solution(n, A, B, array_a):
# 将元素规范化到 [0, 9] 的范围内
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(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]

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)`

运行结果:

image.png