回溯查找元素加起来小于 K 的 n 元向量

Question

我对以下问题感兴趣主要是为了获得对回溯算法的直觉，所以我不是在寻找不使用回溯的替代解决方案。

问题：找到所有n个元素的向量，使得它们的元素之和小于或等于某个数K。向量中的每个元素都是整数。

示例：如果 n = 3，且 K = 10，则 [9, 0, 0] 和 [5, 0, 5] 是解，而 [3, 1, 8] 不是。

从 this site 开始，我改编了 python 代码以尝试实施解决方案。

这是通用的 "backtracking engine" 函数：

def solve(values, safe_up_to, size):

    solution = [None] * size

    def extend_solution(position):
        for value in values:
            solution[position] = value
            if safe_up_to(solution, position):
                if position >= size-1 or extend_solution(position+1):
                    return solution
        return None

    return extend_solution(0)

这里是检查解决方案是否为 "safe so far" 的函数：

def safe_up_to(partial_solution, target = 100): 
   partial_solution = np.array(partial_solution)  # convert to np array 

   # replace None with NaN
   partial_solution = np.where(partial_solution == None, np.nan, partial_solution)

   if np.nansum(partial_solution) <= target: 
       return True
   else: 
       return False

但是，当我运行这两个函数一起使用时，我只得到一个全为零的向量。

solve(values=range(10), safe_up_to=safe_up_to, size=5)

我应该如何修改这段代码以获得所有可行的解决方案？

Answer 1

这是您的代码的轻微修改版本。我试图让它尽可能少地改变：

import numpy as np
from functools import partial

def solve(values, safe_up_to, size):

    solution = [None] * size

    def extend_solution(position):
        for value in values:
            solution[position] = value
            if safe_up_to(solution):
                if position >= size-1:
                    yield np.array(solution)
                else:
                    yield from extend_solution(position+1)
        solution[position] = None

    return extend_solution(0)

def safe_up_to(target, partial_solution): 
   partial_solution = np.array(partial_solution)  # convert to np array 

   # replace None with NaN
   partial_solution = np.where(partial_solution == None, np.nan, partial_solution)

   if np.nansum(partial_solution) <= target: 
       return True
   else: 
       return False 

for sol in solve(values=range(10), safe_up_to=partial(safe_up_to,4), size=2):
    print(sol,sol.sum())

打印：

[0 0] 0
[0 1] 1
[0 2] 2
[0 3] 3
[0 4] 4
[1 0] 1
[1 1] 2
[1 2] 3
[1 3] 4
[2 0] 2
[2 1] 3
[2 2] 4
[3 0] 3
[3 1] 4
[4 0] 4

回溯查找元素加起来小于 K 的 n 元向量

Backtracking to find n-element vectors whose elements add up to less than K

python

algorithm

numpy

backtracking