生成总和为 k 的正整数数组

Question

我的任务很简单：我想生成一个（理想情况下numpy）数组，其中包含m个正数（> = 0）的所有组合，但有界（<= e) integers that sum exactly to k. Note that k and m might be relatively high, so generating all combinations and filtering will not work.

我已经用简单的递归Python实现了它，但是这个小函数占用了我大部分时间，我需要替换它才能更好地执行。我试图想出 numpy/pytorch 代码来生成这个数组，但到目前为止我还没有成功。

我目前在项目中使用 numpy 和 pytorch，但只要我编写 python 代码，我就对其他库开放，并且最终得到可以转换为 numpy 数组的东西。

这是一些代码：

import timeit


def get_summing_up_to(max_degree, sum, length, current=0):
    assert sum >= 0
    assert length >= 1
    if length == 1:
        residual = sum - current
        if residual <= max_degree:
            return [(residual,)]
        else:
            return []
    max_element = min(max_degree, sum - current)
    return [
        (i,) + t
        for i in range(max_element + 1)
        for t in get_summing_up_to(
            max_degree, sum, length - 1,
            current=current + i
        )
    ]


if __name__ == '__main__':
    result = timeit.timeit('get_summing_up_to(60, 60, 6)', globals=globals(), number=1)
    print(f"Execution time: {result} for max_degree=60, sum=60, length=6")

    result = timeit.timeit('get_summing_up_to(30, 30, 8)', globals=globals(), number=1)
    print(f"Execution time: {result} for max_degree=30, sum=30, length=8")

Answer 1

如果你想找到所有组合，你可以使用 itertools。如果您想在每次选择后进行替换，请使用

combinations_with_replacement

，在这种情况下，将允许使用 (m,m) 等组合。否则使用

combination

。

import itertools

e = 10 # the limit
# generate elements to be combined
# so they are from 0 to e
# with 1 spacing in between each
elements = np.linspace(0,e,1)

r = 2 # combinations of r elements
# Generate all combinations of r elements
comb = itertools.combinations_with_replacement(elements, r)

# print all the combinations as tuples
print(list(comb))