迭代具有相同总和的固定大小正整数列表的算法

Question

我正在寻找一种快速且内存效率高的算法来遍历所有可能的具有给定总和 (N) 的相同大小 (S) 的正整数列表。

例如，如果 S = 3 和 N = 4，结果将是（我有一个非常低效的算法）：

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

不一定按这个顺序。另一种看待它的方法是如何将数字 N 切成 S 块。如果我还可以为列表中的每个单独值设置最大值，那么该算法将是完美的。

我将使用它以不同于product(*indices) 生成的顺序遍历多维数组。

同时生成所有索引组合并按总和对其进行排序会太慢/占用内存。

Answer 1

找到了一个解决方案：它基于这样的想法，即正数 N 是一行单位，将它们分成 S 块是在列表中放置 (S-1) 个分隔符的问题。 这些分隔符可以使用combinations(range(N + S - 1), S - 1) 进行迭代。下一步是计算分隔符之前、之间和之后的单元数：

def partition(N, size):
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        yield [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]

当您想限制结果中的每个项目时，您可以过滤掉不需要的项目（当然不是最佳的，但我想使用combinations，因为它可能是用 C 实现的，因此可能比我想出的任何东西都快得多在 python 中）。简单版：

def sized_partition_slow(N, sizes):
    size = len(sizes)
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        result = [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]
        if all(r < s for r, s in zip(result, sizes)):
            yield result

还有更快但更复杂的版本：

def sized_partition(N, sizes):
    size = len(sizes)
    n = N + size - 1
    for splits in combinations(range(n), size - 1):
        result = []
        for s, s0, s1 in zip(sizes, (-1,) + splits, splits + (n,)):
            r = s1 - s0 - 1
            if r >= s:
                break
            result.append(r)
        else:
            yield result

我用这个作为早期测试：

for indices in partition(4, 3):
    assert sum(indices) == 4
    assert all(0 <= i for i in indices)

for indices in sized_partition(4, [3, 3, 3]):
    assert sum(indices) == 4
    assert all(0 <= i < 3 for i in indices)

顺便说一句：从臀部：您可以通过迭代 S（大小）来生成整数分区问题的解决方案：如：

def integer_partition(N, order=False):
    result = set()
    for size in range(1, N+1):
        for splits in combinations(range(1, N), size - 1):
            if order:
                p = tuple(s1 - s0 for s0, s1 in zip((0,) + splits, splits + (N,)))
            else:
                p = tuple(sorted(s1 - s0 for s0, s1 in zip((0,) + splits, splits + (N,))))
            result.add(p)
    return sorted(result, key=lambda r: (len(r), r))

我稍微调整了combinations() 迭代器以不给零。如果order=False，它将为具有不同顺序的相同分区减倍。

Answer 2

这是我认为最简单的方法：

def elect(S,N,List):
    result_list = []
    for list_val in List:
        if sum(list_val) == N:
            if len(list_val) == S:
                result_list.append(list_val)

    return result_list

这对 100 万个列表在 1 秒内有效。如果你想加快速度，你可以加入其他 if 语句，例如 if sum(list_val[0:N/2]) > N 或 len(list_val) / 2 > S：这样的语句可以更快地检测到情况。

另一种方法是对列表进行排序并查看第一个 N 数字总和。如果它大于您想要的，您可以选择这些列表。