刷题笔记-二分数字组合 | 豆包MarsCodeAI刷题

107 阅读9分钟

二分数字组合

问题描述

我们需要将一个整数数组 array_a 划分为两组,使得每组的和的个位数分别等于给定的 AB。允许的特殊情况是其中一组可以为空。任务是计算所有可能的划分方式。

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

image.png

问题分析

原始代码的解题思路主要基于回溯算法。此方法通过递归尝试将每个元素放入两个不同的组中,每一步都维护两个变量 sum1sum2 来记录当前两组的和。在递归的最后一层,当所有元素都被处理完后,会检查这两组和的个位数是否符合目标 AB 的要求。如果符合,就增加计数器。

此外,为了处理特殊的边界条件,代码还需要检查数组的总和的个位数是否直接等于 AB,以确认是否可能将一组视为空。整体思路清晰,但由于使用了回溯方法,可能会导致重复计算,从而影响效率。

思路转换代码

1. 问题定义

  • 输入:整数数组 array_a 和两个目标个位数 AB
  • 输出:满足条件的划分方式的数量。

2. 主要思路

  • 使用回溯算法递归地尝试将每个数字放入两个组。
  • 使用 sum1sum2 来记录当前两组的和。
  • 当遍历完所有元素后,检查和的个位数是否符合要求。

3. 特殊情况处理

  • 在计算完所有可能的划分后,检查数组总和的个位数是否等于 AB,这意味着可以将一组视为空。

代码实现

def solution(n, A, B, array_a):
    total_sum = sum(array_a)  # 计算数组的总和
    count = 0  # 初始化计数器,用于记录符合条件的划分方式

    # 回溯函数,用于计算所有可能的划分方式
    def backtrack(index, sum1, sum2):
        nonlocal count  # 允许在嵌套函数中修改外部变量 count
        if index == n:  # 如果已经处理完所有元素
            if (sum1 % 10 == A and sum2 % 10 == B) or (sum1 % 10 == B and sum2 % 10 == A):
                count += 1  # 如果当前两组的和的个位数符合要求,计数器加1
            return

        # 尝试将当前索引的元素加入第一组
        backtrack(index + 1, sum1 + array_a[index], sum2)
        # 尝试将当前索引的元素加入第二组
        backtrack(index + 1, sum1, sum2 + array_a[index])

    # 从第一个数字开始回溯
    backtrack(0, 0, 0)

    # 检查特殊情况:一组为空
    if total_sum % 10 == A or total_sum % 10 == B:
        count += 1  # 如果总和的个位数等于 A 或 B,计数器加1

    return count  # 返回符合条件的划分方式数量

具体实现步骤

  1. 初始化和定义回溯函数

    • 计算数组的总和 total_sum
    • 初始化计数器 count,用于记录符合条件的划分方式数量。
    • 定义 backtrack 函数,来探索所有可能的划分。
  2. 递归回溯

    • 在回溯过程中,检查是否已经遍历完所有元素(即 index == n)。
    • 如果是,检查和的个位数是否符合目标条件,若符合则增加 count
  3. 特殊情况处理

    • 回溯计算结束后,检查整体数组总和的个位数,可以直接判断是否能形成满足条件的划分。
  4. 返回结果

    • 返回计数器 count,即符合条件的划分方式数量。

示例和测试

在实现后,代码提供了一些测试用例来验证功能。通过这些示例,可以确认输出结果是否符合预期。

if __name__ == "__main__":
    # 测试样例
    print(solution(3, 1, 2, [1, 1, 1]))  # 预期输出:3, 实际输出:6
    print(solution(3, 3, 5, [1, 1, 1]))  # 预期输出:1
    print(solution(2, 1, 1, [1, 1]))    # 预期输出:2
    print(solution(5, 3, 7, [2, 3, 5, 7, 9]))  # 预期输出:0
    print(solution(13, 8, 3, [21, 9, 16, 7, 9, 19, 8, 4, 1, 17, 1, 10, 16]))  # 预期输出:1

问题:代码逻辑的缺陷

在原始代码中,特定输入(例如 print(solution(3, 1, 2, [1, 1, 1])))的实际输出与预期不符,可能是因为在回溯过程中存在重复计数的情况。相同的元素组合可能因为不同的加入顺序被多次计算,这导致了计数的错误。

代码优化

优化后的代码使用动态规划来避免回溯中的重复计算,显著提高效率:

def solution(n, A, B, array_a):
    array_a = [x % 10 for x in array_a]  # 只保留个位数
    
    total_sum = sum(array_a)  # 计算总和
    total_sum %= 10  # 只保留总和的个位数
    if (total_sum == A or total_sum == B):
        return 1  # 如果总和的个位数等于 A 或 B,返回1
    if (total_sum != (A + B) % 10):
        return 0  # 如果总和的个位数不等于 A 和 B 的和的个位数,返回0

    # dp数组,f[i][j]表示前i个数字中和的个位数为j的方式数量
    f = [[0 for _ in range(10)] for _ in range(n + 1)]
    f[0][0] = 1  # 初始化:选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]  # 返回前n个数字中和的个位数为B的方式数量

代码重要步骤及含义

  1. 处理数组元素
    • 将输入数组 array_a 中每个元素取模10,保留其个位数。这是因为我们只关注数字的个位数,减少了计算的复杂性。
    array_a = [x % 10 for x in array_a]  # 只保留个位数
  1. 计算总和的个位数
    • 计算数组中所有元素的总和,然后取模10,得到总和的个位数。这一步是为了后续判断是否可以满足分组条件。
    total_sum = sum(array_a)  # 计算总和
    total_sum %= 10  # 只保留总和的个位数
  1. 处理特殊情况
    • 如果总和的个位数直接等于 AB,那么我们可以将一组视为空,返回1。这个返回值代表满足条件的一个有效划分。
    if (total_sum == A or total_sum == B):
        return 1  # 如果总和的个位数等于 A 或 B,返回1
  1. 检查可行性
    • 如果总和的个位数不等于 AB 的和的个位数,那么不可能有有效的划分,直接返回0。
    if (total_sum != (A + B) % 10):
        return 0  # 如果总和的个位数不等于 A 和 B 的和的个位数,返回0
  1. 初始化动态规划数组
    • 创建一个二维数组 f,其中 f[i][j] 表示前 i 个数字中,和的个位数为 j 的划分方式数量。f[0][0] = 1 表示选取0个数字,和为0的唯一一种方式。
    f = [[0 for _ in range(10)] for _ in range(n + 1)]
    f[0][0] = 1  # 初始化:选0个数字,和为0的方式数量为1
  1. 动态规划填表
    • 外层循环遍历每个数字,内层循环遍历和的个位数(0-9)。
    • f[i][j] += f[i - 1][j] 表示当前数字不被选择,保持当前和的位数 j 的计数。
    • f[i][j] += f[i - 1][(j - array_a[i - 1] + 10) % 10] 表示当前数字被选择,更新和的个位数,确保结果在0-9的范围内。
    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]  # 选择当前数字
  1. 返回结果
    • 最后,返回 f[n][B],即使用所有 n 个数字来形成和的个位数为 B 的方式数量。这是我们需要的划分结果。
    return f[n][B]  # 返回前n个数字中和的个位数为B的方式数量

两版代码对比

特征原始代码优化后代码
计算方式使用回溯算法,递归尝试所有可能的划分方式使用动态规划,减少重复计算,提升效率
状态存储变量 sum1sum2 存储当前两组的和使用二维数组 f 存储不同情况下的和的计数
复杂度时间复杂度较高,最坏情况下为 O(2^n)时间复杂度为 O(n),空间复杂度 O(n)
特殊情况处理处理一组为空的情况,需手动检查总和的个位数直接在总和计算时处理,逻辑更清晰
直接返回的情况需要遍历所有情况后,进行最终判断通过条件直接返回结果,提高效率

问题关键点总结

关键点描述
目标将数组划分为两组,使得各组和的个位数分别为 A 和 B 或 A 和 B 互换。
允许特殊情况可以允许其中一组为空,要求剩余数字的和的个位数为 A 或 B。
计算方式优化后的代码使用动态规划,避免了重复计算,提升效率。
状态转移使用 f[i][j] 存储前 i 个数字中和的个位数为 j 的方式数量。
边界条件通过计算总和的个位数,直接判断是否满足条件,简化逻辑。

问题总结

通过这个练习,我们不仅学习了如何将数组划分为两组的问题解决方法,还理解了动态规划的应用。在原始回溯方法中,虽然能够解决问题,但由于存在重复计算和复杂的状态管理,导致效率低下。优化后的动态规划方法,利用了状态转移的思想,大大减少了时间复杂度,使得代码更简洁,并能够处理更大的输入数据。

在实际编程中,常常需要权衡不同算法的优缺点,选择合适的算法来应对特定问题。在此过程中,深入理解问题的背景和条件是至关重要的。希望这个刷题笔记能够帮助你在未来的练习中更好地理解和应用算法。