拆分总和 K 方式组合数学 python

Question

我有一个总密度 N。我需要找到所有可以将该密度分成 k 组的方法，假设一些最小可分部分 d。

所以，如果 N=4，k=3，d=1，我需要：

[
    (1, 1, 2),
    (1, 2, 1),
    (2, 1, 1),
]

如果 N=5，k=3，d=1：

[
    (1, 1, 3),
    (1, 3, 1),
    (3, 1, 1),
    (1, 2, 2),
    (2, 1, 2),
    (2, 2, 1),
]

感觉这应该符合一些基本的组合问题，但我想不出来。 这篇文章在 R 中引用了一个非常相似的操作，但不是 python。

一种天真的做事方式是将itertools.product的结果过滤为仅总和为 N 的元组，但这感觉不雅：

import itertools


def compositions(N, k, d=1):
    for seq in itertools.product(*[range(d, N, d)] * k):
        if sum(seq) == N:
            yield seq

例子

>>> list(compositions(7, 3))
[(1, 1, 5),
 (1, 2, 4),
 (1, 3, 3),
 (1, 4, 2),
 (1, 5, 1),
 (2, 1, 4),
 (2, 2, 3),
 (2, 3, 2),
 (2, 4, 1),
 (3, 1, 3),
 (3, 2, 2),
 (3, 3, 1),
 (4, 1, 2),
 (4, 2, 1),
 (5, 1, 1)]

还有其他想法吗？

Answer 1

您可以假设 d=1，因为如果不是，那么您可以使用 d 对 N 进行归一化，即 N' = N/d，现在 d=1。

代码如下所示：

import itertools

def divisions(N, k):
    for q in itertools.combinations(range(N - 1), k - 1):
        yield [j - i for i, j in zip((-1,) + q, q + (N - 1,))]

answer = [x for x in divisions(7, 3)]
print(answer)