二分数字组合题解

简单分享一下"二分数字组合"这个题目的解题方法:

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

思路: 首先我们可以求出所有数字的和, 记作 $sum$ 如果 (A + B) % 10 != sum % 10, 那么一定不存在两个个位为A 和 B 的数相加为sum, 所以只需要考虑第三种情况就可以了

设两个集和为 C集和 和 D集和 否则, 如果C集和的尾巴是A, 那么D集和的尾巴一定是B, C集和的尾巴是B, 那么D集和的尾巴一定是A

我们使用dp来求方法数, $dp[i][j]$ 表示已经处理了 (i) 个数, 其中属于C集和数字和的个位为 j

dp转移的时候, 我们对于 a[i] 有两种处理方法:

1. 将 a[i] 加入集和C, 此时个位数发生改变, 将dp[i][j] 加到 dp[i + 1][(j + a[i]) % 10] 上
2. 将 a[i] 加入集和D, 此时个位数没有变动, 将dp[i][j] 加到 dp[i + 1][j] 上
3. 最终答案就是 dp[n][A] + dp[n][B]
4. 复杂度 10n
5. 写得可能不是很好, 欢迎大家指正

def solution(n, A, B, array_a):
    sum = 0
    for i in range(n):
        sum += array_a[i]
    
    if (A + B) % 10 != sum % 10:
        # 只需要考虑第三种情况
        ans = 0
        if(sum % 10 == A or sum % 10 == B):
            ans += 1
        return ans
        
    dp = [[0 for i in range(10)] for i in range(n + 1)] 
    dp[1][0] = 1
    for i in range(n):
        for t in range(10):
            aim = (t + array_a[i]) % 10
            dp[i + 1][t] += dp[i][t]   #将a[i] 加入B集和
            dp[i + 1][aim] += dp[i][t] #将a[i] 加入A集和
            
    ans = dp[n][A] + dp[n][B]
    return ans