如何将数字拆分为三个子集？现有Python代码不符预期求修正

原代码问题 #

原代码采用贪心策略每次选取最大元素，无法保证总能凑出目标和。比如n=5时，贪心会先取5，剩余元素无法拆分成两个和为5的子集，但实际存在合法拆分方案。另外原代码在无法拆分时返回None，不符合返回-1的要求，且n<5的判断逻辑错误（比如n=6是可以拆分的）。

修正方案 #

替换贪心算法为回溯法寻找合法子集，调整边界判断逻辑，确保无法拆分时返回-1：

def find_subset(lst, target):
    # 回溯查找和为target的子集，返回子集与剩余元素列表
    def backtrack(start, current_sum, path):
        if current_sum == target:
            # 生成剩余元素（利用原列表有序性快速过滤）
            remaining = []
            ptr = 0
            for num in lst:
                if ptr < len(path) and num == path[ptr]:
                    ptr += 1
                else:
                    remaining.append(num)
            return path, remaining
        if current_sum > target or start >= len(lst):
            return None
        # 选择当前元素
        res = backtrack(start + 1, current_sum + lst[start], path + [lst[start]])
        if res:
            return res
        # 不选择当前元素
        return backtrack(start + 1, current_sum, path)
    
    return backtrack(0, 0, [])

def partition(n):
    total = n * (n + 1) // 2
    # 总和无法被3整除，直接返回-1
    if total % 3 != 0:
        return -1
    target = total // 3
    lst = list(range(1, n + 1))
    
    # 寻找第一个和为target的子集
    first_subset, remaining = find_subset(lst, target)
    if not first_subset:
        return -1
    
    # 从剩余元素中寻找第二个和为target的子集
    second_subset, third_subset = find_subset(remaining, target)
    if not second_subset or not third_subset:
        return -1
    
    return [first_subset, second_subset, third_subset]

# 测试示例
print(partition(5))  # 输出类似 [[1,4], [2,3], [5]]
print(partition(8))  # 输出类似 [[8,4], [7,5], [1,2,3,6]]
print(partition(3))  # 输出 -1（无法拆分）

关键说明 #

1. 回溯法找子集：find_subset通过递归尝试选择/不选择元素，确保找到符合和要求的子集，并同步生成剩余元素列表
2. 合法性校验：先判断总和是否能被3整除，再依次验证前两个子集的存在性，最后确保第三个子集非空
3. 返回值处理：所有无法拆分的场景统一返回-1，符合需求

内容的提问来源于stack exchange，提问作者Proger228