Python求解互异整数立方和匹配目标值的代码修正问题

需求为生成元素互不相同的整数列表，要求列表内所有整数的立方和恰好等于给定目标数。现有实现代码在多数测试用例下可正常运行，但计算如下目标值时输出不符合预期：

n=sum([n*n*n for n in range(1001)]) 

该场景下预期正确结果为 [6303, 457, 75, 14, 9, 7, 5, 4]，程序实际输出为 [6303, 457, 75, 15, 8, 4, 3, 2, 1]。

def sum_of_cubes(n):
    original=n
    i=1
    lst=[]
    tot=0
    while i**3<=n:
        i+=1
    lst.append(i-1)
    n-=(i-1)**3
    for j in range(lst[0],0,-1):
        if j**3<n:
            lst.append(j)
            n-=j**3
    if n==1:
        lst.append(n)
    for i in lst:
        tot+=i**3
    #if original-tot>1:
        #return None
    return lst

n=sum([n*n*n for n in range(1001)])
print(sum_of_cubes(n))

现有代码采用最基础的贪心逻辑：从大到小遍历整数，只要当前数的立方小于剩余目标值就直接加入结果列表。这种只选局部最优的策略无法保证得到全局正确解：在这个测试用例里，选15之后剩余的差值只能靠多个小数字凑和，但跳过15选14就能凑出正确组合，单次贪心没有回溯纠错的能力，自然输出错误结果。

替换单次贪心逻辑为带剪枝的深度优先回溯搜索即可：

• 每次从大到小枚举可选整数，优先选大数加快搜索速度，同时限制后续选的数必须比当前选的数小，保证列表元素互不重复
• 加剪枝判断：如果当前剩余需要凑的目标值，比当前可选的所有更小数字的立方总和还大，直接跳过这个分支，不用做无效遍历
• 一旦找到刚好凑够立方和的组合就立刻返回，不用遍历所有可能性

修正后可运行代码如下：

def find_cube_sum(target):
    # 预计算立方值避免重复运算
    max_num = int(target ** (1/3)) + 2
    cube = [i**3 for i in range(max_num + 1)]
    # 预计算前缀立方和用于剪枝
    prefix_sum = [0]*(max_num + 1)
    for i in range(1, max_num+1):
        prefix_sum[i] = prefix_sum[i-1] + cube[i]
    
    def dfs(remain, upper_bound, path):
        if remain == 0:
            return path.copy()
        if remain < 0 or upper_bound == 0:
            return None
        # 剩余所有数加起来都不够凑目标，直接剪枝
        if prefix_sum[upper_bound] < remain:
            return None
        # 从大到小枚举提速
        for num in range(upper_bound, 0, -1):
            if cube[num] > remain:
                continue
            path.append(num)
            res = dfs(remain - cube[num], num-1, path)
            if res is not None:
                return res
            path.pop()
        return None
    
    return dfs(target, max_num, [])

n = sum([i*i*i for i in range(1001)])
print(find_cube_sum(n))

运行代码即可输出预期的正确结果。

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