HJ59 数组分组

题目描述

对于给定的 n 个整数，将其分为 a、b 两个数组，满足：

• 所有 5 的倍数元素均在 a 数组中
• 所有 3 的倍数元素（不包括 5 的倍数）均在 b 数组中
• 其他元素可以任意分配

求解是否存在一种分配方案，使得 a 数组中各个元素之和等于 b 数组中各个元素之和。

解题思路

这道题需要用到深度优先搜索(DFS) + 剪枝优化。

核心思路

1. 预处理：将元素按规则分为三类
   • 5 的倍数 → 必须放入 a 数组
   • 3 的倍数（非 5 的倍数）→ 必须放入 b 数组
   • 其他 → 可自由分配
2. DFS 搜索：对于可自由分配的元素，尝试两种选择
   • 加入 a 数组
   • 加入 b 数组
3. 剪枝优化：使用索引 + 当前和的方式传参，避免重复计算

代码实现

n = int(input())
nums = list(map(int, input().split()))

a = []   # 5的倍数，必须放a组
b = []   # 3的倍数，必须放b组
rem = [] # 其他，可任意分配

for x in nums:
    if x % 5 == 0:
        a.append(x)
    elif x % 3 == 0:
        b.append(x)
    else:
        rem.append(x)

sum_a = sum(a)  # a组的初始和
sum_b = sum(b)  # b组的初始和

def dfs(i, sa, sb):
    if i == len(rem):
        return sa == sb
    # 尝试将 rem[i] 加到 a 组或 b 组
    return dfs(i + 1, sa + rem[i], sb) or dfs(i + 1, sa, sb + rem[i])

ans = dfs(0, sum_a, sum_b)
print("true" if ans else "false")

代码解析

核心函数 dfs(i, sa, sb)

参数	含义
i	当前考虑的可自由分配元素的索引
sa	a 数组当前的累加和
sb	b 数组当前的累加和

算法流程

分类阶段：
┌─────────┬─────────────┬─────────────┐
│ 5的倍数 │  3的倍数    │   其他      │
│  → a组  │  → b组      │ → DFS分配   │
└─────────┴─────────────┴─────────────┘

DFS阶段：
dfs(0, sum_a, sum_b)
    ├── 加入rem[0]到a组 → dfs(1, sa+rem[0], sb)
    └── 加入rem[0]到b组 → dfs(1, sa, sb+rem[0])

关键点

1. 注意优先级：先判断 5 的倍数，再判断 3 的倍数（因为同时是 3 和 5 的倍数即 15 的倍数应归入 a 组）
2. 剪枝：传入当前和而非整个数组，避免重复求和
3. 短路求值：使用 or 连接两个递归调用，找到解后立即返回

示例演示

示例 1：

输入: 4
     1 5 -5 1

分类: a=[5,-5], sum_a=0
     b=[], sum_b=0
     rem=[1,1]

DFS过程:
dfs(0, 0, 0) → dfs(1, 1, 0) → dfs(2, 2, 0) → 2≠0 ✗
                          → dfs(2, 1, 1) → 1=1 ✓

输出: true

示例 2：

输入: 3
     3 5 8

分类: a=[5], sum_a=5
     b=[3], sum_b=3
     rem=[8]

DFS过程:
dfs(0, 5, 3) → dfs(1, 13, 3) → 13≠3 ✗
            → dfs(1, 5, 11) → 5≠11 ✗

输出: false

复杂度分析

• 时间复杂度：O(2^k) - k 为可自由分配元素的个数，最多 30 个
• 空间复杂度：O(k) - 递归栈深度

优化建议

1. 提前剪枝：如果初始 sum_a 和 sum_b 之差超过剩余元素总和，直接返回 false
2. 记忆化：用 memo[(i, diff)] 记录已访问状态，避免重复计算
3. 位运算：位掩码方式记录状态，进一步优化空间

由于题目中 n ≤ 30 且数据随机生成，基础 DFS 已可通过测试。