91087


和がkとなるようなn個の非負整数の順列を重複なく全列挙する方法が知りたいです．

ナイーブに実装しましたが`O(k^n)`，これより計算量が改善される解法は存在しますか．

[code]

    # Python 3.10.5
    
    from itertools import product
    
    def sol(n, k):
        for perm in product(range(k+1), repeat=n):
            if sum(perm) == k:
                yield perm
    
    print(*sol(n=2, k=4))
    # (0, 4) (1, 3) (2, 2) (3, 1) (4, 0)
    print(*sol(n=3, k=3))
    # (0, 0, 3) (0, 1, 2) (0, 2, 1) (0, 3, 0) (1, 0, 2) (1, 1, 1) (1, 2, 0) (2, 0, 1) (2, 1, 0) (3, 0, 0)
    
[/code]


CC BY-SA 4.0
2022-09-13T14:41:48.447
91086
2022-09-13T18:19:31.520
51374
question
アルゴリズム
和がkとなるようなn個の非負整数の順列を全列挙する


数値の分割位置の組み合わせで考えればよさそうです。

k個の「〇」が並べてあって、「〇」の間に間仕切り「|」を(n-1)個入れて分割することを考えます。  
例えば、k=n=3の場合は、以下のようになります。

||〇〇〇 = (0,0,3)  
|〇|〇〇 = (0,1,2)  
...  
〇〇〇|| = (3,0,0)

さて、間仕切りを配置できる箇所は、「〇」の間と前後で(k+1)箇所あります。  
間仕切りは、同じ場所に複数配置できるので、重複組み合わせを使えばパターンを網羅できます。  
重複組み合わせは、Pythonだと`itertools`の`combinations_with_replacement`で生成できます。  
あとは、間仕切りの配置場所から、数値の組み合わせへの変換を行えば、目的の順列の全列挙ができます。

[code]

    from itertools import combinations_with_replacement, pairwise
    
    def sol(n, k):
        for c in combinations_with_replacement(range(k + 1), n - 1):
            yield tuple(b - a for a, b in pairwise((0, *c, k)))
    
[/code]


CC BY-SA 4.0
2022-09-13T15:48:17.613
91087
2022-09-13T18:19:31.520
2022-09-13T18:19:31.520
43506
43506
91086
answer


若干計算量が改善された解法を共有します．

[code]

    # Python 3.10.5
    
    from itertools import product
    
    def sol(n, k):
        for perm in product(range(k+1), repeat=n-1):
            if (x := sum(perm)) <= k:
                yield *perm, k-x
    
    
[/code]


CC BY-SA 4.0
2022-09-13T17:04:50.257
91091
2022-09-13T17:04:50.257
51374
91086
answer