python - 非负整数组合的高效枚举

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]] .

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

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)] .

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

编辑

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

在左边，我们有 l=2*n在右边我们有l=n+1 .在这两种情况下，当 n<=5 时，user2357112 的第二个解决方案比其他解决方案更快。 .对于 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_with_replacement 的方法:从 l 个零的列表开始，从 中选择 n 个位置进行替换l 可能的位置，并对每个选择的位置加 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