非负整数组成的高效枚举

Question

我想写一个函数my_func(n,l)，对于一些正整数n，有效地枚举长度为l的有序非负整数组合*（其中l大于n） .比如我想让my_func(2,3)返回[[0,0,2],[0,2,0],[2,0,0],[1,1,0],[1,0,1],[0,1,1]]。

我最初的想法是将现有代码用于正整数分区（例如 this post 中的 accel_asc()），将正整数分区扩展几个零并返回所有排列。

def my_func(n, l):
    for ip in accel_asc(n):
        nic = numpy.zeros(l, dtype=int)
        nic[:len(ip)] = ip
        for p in itertools.permutations(nic):
            yield p

这个函数的输出是错误的，因为每个数字出现两次（或多次）的非负整数组合在my_func的输出中出现了多次。例如，list(my_func(2,3)) 返回[(1, 1, 0), (1, 0, 1), (1, 1, 0), (1, 0, 1), (0, 1, 1), (0, 1, 1), (2, 0, 0), (2, 0, 0), (0, 2, 0), (0, 0, 2), (0, 2, 0), (0, 0, 2)]。

我可以通过生成所有非负整数组合的列表、删除重复条目、然后返回剩余列表（而不是生成器）来纠正此问题。但这似乎非常低效，并且可能会遇到内存问题。有什么更好的方法来解决这个问题？

编辑

我对这篇文章和 cglacet 在 cmets 中指出的 another post 的答案中提供的解决方案进行了快速比较。

左边是l=2*n，右边是l=n+1。在这两种情况下，user2357112 的第二个解决方案比其他解决方案更快，n<=5。对于n>5，user2357112、Nathan Verzemnieks 和 AndyP 提出的解决方案或多或少是相关的。但在考虑l 和n 之间的其他关系时，结论可能会有所不同。

.......

*我最初要求的是非负整数分区。 Joseph Wood 正确地指出，我实际上是在寻找整数组合，因为序列中数字的顺序对我很重要。

Answer 1

使用stars and bars 概念：选择在n 星之间放置l-1 条的位置，并计算每个部分中有多少星：

import itertools

def diff(seq):
    return [seq[i+1] - seq[i] for i in range(len(seq)-1)]

def generator(n, l):
    for combination in itertools.combinations_with_replacement(range(n+1), l-1):
        yield [combination[0]] + diff(combination) + [n-combination[-1]]

我在这里使用了combinations_with_replacement 而不是combinations，因此索引处理与combinations 所需要的有点不同。带有combinations 的代码更符合标准的星形和条形处理方式。

---

另外，使用combinations_with_replacement 的另一种方式：从l 零的列表开始，从l 可能的位置中选择n 位置并替换，然后将每个选定位置加1 以产生输出：

def generator2(n, l):
    for combination in itertools.combinations_with_replacement(range(l), n):
        output = [0]*l
        for i in combination:
            output[i] += 1
        yield output

Answer 2

从一个简单的递归解决方案开始，它与您的问题相同：

def nn_partitions(n, l):
    if n == 0:
        yield [0] * l
    else:
        for part in nn_partitions(n - 1, l):
            for i in range(l):
                new = list(part)
                new[i] += 1
                yield new

也就是说，对于下一个较小数字的每个分区，对于该分区中的每个位置，将该位置的元素加 1。它会产生与您相同的副本。不过，我记得一个类似问题的技巧：当您将n 的分区p 更改为n+1 的一个分区时，将p 的所有元素修复到您增加的元素的左侧。也就是说，跟踪p 的修改位置，并且永远不要修改p 左侧的任何“后代”。这是代码：

def _nn_partitions(n, l):
    if n == 0:
        yield [0] * l, 0
    else:
        for part, start in _nn_partitions(n - 1, l):
            for i in range(start, l):
                new = list(part)
                new[i] += 1
                yield new, i

def nn_partitions(n, l):
    for part, _ in _nn_partitions(n, l):
        yield part

非常相似——只是在每一步都传递了额外的参数，所以我添加了包装器来为调用者移除它。

我尚未对其进行广泛测试，但这似乎相当快 - nn_partitions(3, 5) 大约 35 微秒，nn_partitions(10, 20) 大约 18 秒（产生超过 2000 万个分区）。 （来自user2357112 的非常优雅的解决方案对于较小的情况大约需要两倍的时间，对于较大的情况大约需要四倍的时间。编辑：这是指该答案中的第一个解决方案；在某些情况下，第二个解决方案比我的要快在其他情况下速度较慢。）