python - Python 中的组合函数

Question

我真的很难尝试编写一个函数组合，它以两个自然数 k 和 n 作为输入并返回所有大小为 k 且总和为 n 的元组的集合。我关心的是找到同一组数字的不同排列。

我从 python 中找到了这个文档，可以在不使用 itertools 的情况下进行计算。

我的问题允许使用这些库

import sys
import numpy as np
import scipy as sp
from scipy.special import *


def compositions(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = range(r)
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

Given input compositions(3, 4)
output:
all possible combinations
1 + 2 + 1
2 + 1 + 1
1 + 1 + 2

actual output from composition(3,4):
{(1, 1, 2,), (1, 2, 1), (2, 1, 1)}

如有任何帮助，我们将不胜感激。

Answer 1

首先，请注意随着 n，解决方案的长度增长得非常快，因此如果您使用大量数字，计算这个的时间和内存可能会大打折扣。

话虽如此，但请注意，获取所有数字组合并检查总和并不是一个好主意，因为您会生成很多情况，只需检查第一个数字就知道它们不会起作用。换句话说，您需要尽快修剪数组的搜索。

思考这类问题和更好更快的实现是非常有趣的。这里有一种可能性:

def gen_combinations(k, n):
    assert n > k > 1
    to_process = [[i] for i in range(1, n+1)]
    while to_process:
        l = to_process.pop()
        s = sum(l)
        le = len(l)
        #If you do not distiguish permutations put xrange(l[-1],n-s+1) 
        # instead, to avoid generating repeated cases.
        # And if you do not want number repetitions, putting
        # xrange(l[-1] + 1, n-s+1) will do the trick.
        for i in xrange(1, n-s+1): 
            news = s + i
            if news <= n:
                newl = list(l)
                newl.append(i)
                if le == k-1 and news == n:
                    yield tuple(newl)
                elif le < k-1 and news < n:
                    to_process.append(newl)

这是一个生成器，如果您需要列表，只需执行，例如 list(gen_combinations(3,4))。