MarsCode-二分数字组合

二分数字组合

问题描述

小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

题解

解法一 DFS

def solution(n, A, B, array_a):    
    # 处理特殊情况：如果一组可以为空
    total_sum = sum(array_a) % 10
    if total_sum == A or total_sum == B:
        result = 1
    else:
        result = 0
    
    # 使用 DFS 处理一般情况
    def dfs(index, sumA, sumB):
        # 如果已经处理到末尾
        if index == n:
            if sumA % 10 == A and sumB % 10 == B:
                return 1
            return 0
        
        current_number = array_a[index]

        return dfs(index + 1, sumA + current_number, sumB)+dfs(index + 1, sumA, sumB + current_number)

    # 加上一般情况的结果
    result += dfs(0, 0, 0)

    return result

# 测试用例
if __name__ == "__main__":
    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)
    print(solution(5, 3, 7, [2, 3, 5, 7, 9]) == 0)

时间复杂度$O(2^n)$，出现超时

解法二 动态规划

def solution(n, A, B, array_a):
    # Please write your code here
    # 先处理特殊情况，即一组为空
    total_sum = sum(array_a)
    total_sum_mod = total_sum % 10
    if total_sum_mod == A or total_sum_mod == B:
        result = 1
    else:
        result = 0

    # 初始化动态规划表，dp[i][modA][modB] 表示前 i 个元素，A 和 B 组个位数分别为 modA 和 modB 的方式数
    dp = [[[0] * 10 for _ in range(10)] for _ in range(n + 1)]
    dp[0][0][0] = 1  # 初始化条件：没有元素时，和为 0
    
    for i in range(1, n + 1):
        num = array_a[i - 1] % 10  # 当前数字取个位数
        for modA in range(10):
            for modB in range(10):
                if dp[i - 1][modA][modB] > 0:
                    # 选择将当前数字加入到 A 组
                    new_modA = (modA + num) % 10
                    dp[i][new_modA][modB] += dp[i - 1][modA][modB]
                    
                    # 选择将当前数字加入到 B 组
                    new_modB = (modB + num) % 10
                    dp[i][modA][new_modB] += dp[i - 1][modA][modB]
                    
    result += dp[n][A][B]
    return result

时间复杂度：O(n * 10 * 10)，因为我们需要遍历每个数字，并在 10 个 modA 和 10 个 modB 状态中更新。

空间复杂度：O(n * 10 * 10)，存储动态规划表。