python - 检查列表中元素的线性组合是否等于某个和的函数

Question

问题是创建一个函数来检查该列表中某些项目的线性组合加起来是否等于某个总和。结果将是一个包含元组的列表(与列表的长度相同)。例如:给定列表:[3,7,10]，sum= 60结果:[(0, 0, 6), (1, 1, 5), (2, 2, 4), (3, 3, 3), (4, 4, 2), (5, 5 , 1), (6, 6, 0), (10, 0, 3), (11, 1, 2), (12, 2, 1), (13, 3, 0), (20, 0, 0 )]问题是列表的长度是变化的。我尝试用一​​堆 if 语句而不是使用 for 循环来解决它，但必须有更有效的方法来做到这一点。

这是我使用的一些代码。

def get_index(l, s):
    res = []
    if len(l)==3:
        for i in range(s+1):
                for j in range(s+1):
                    for k in range(s+1):
                        if l[0]*i + l[1]*j + l[2]*k==s:
                            res.append((i,j,k))
    return res

已经谢谢了!!

注意:如果我确实将范围更改为 (s//l[i]+1)，它会起作用。

Answer 1

我觉得有更好的方法来做到这一点，但这里有一个使用数组的粗暴方法:

A = np.array([3,7,10])
b = np.array([60])

from itertools import product
combin = [np.arange(i) for i in (b//A)+1]
d = np.stack(list(product(*combin)))
[tuple(i) for i in d[d@A==b]]

或等效地没有 itertools:

d = np.rollaxis(np.indices((b//A)+1),0,4).reshape(-1,3)
[tuple(i) for i in d[d@A==b]]

输出:

[(0, 0, 6), (1, 1, 5), (2, 2, 4), (3, 3, 3), (4, 4, 2), (5, 5, 1), (6, 6, 0), (10, 0, 3), (11, 1, 2), (12, 2, 1), (13, 3, 0), (20, 0, 0)]

---

比较:

#@Ehsan's solution 1
def m1(b):
  combin = [np.arange(i) for i in (b//A)+1]
  d = np.stack(list(product(*combin)))
  return [tuple(i) for i in d[d@A==b]]

#@Ehsan's solution 2
def m2(b):
  d = np.rollaxis(np.indices((b//A)+1), 0, 4).reshape(-1,3)
  return [tuple(i) for i in d[d@A==b]]

#@mathfux solution
def m3(b):
  A, B, C = range(0, b+1, 3), range(0, b+1, 7), range(0, b+1, 10)
  triplets = list(product(A, B, C)) #all triplets
  suitable_triplets = list(filter(lambda x: sum(x)==b, triplets)) #triplets that adds up to 60
  return [(a//3, b//7, c//10) for a, b, c in suitable_triplets]

性能:

in_ = [np.array([n]) for n in [10,100,1000]]

这使得 m2 成为其中最快的。