Function for compositions in Python

时间:2018-03-09 19:09:36

标签: python statistics

I am having a really difficult time trying to Write a function compositions that takes two natural numbers k and n as inputs and returns the set of all tuples of size k that sum to n. I am concerned with finding the different permutations for the same set of numbers.

I have found this documentation from python to calculate without using itertools.

My questions allows for these libraries to be used

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)}

Any help would be appreciated.

2 个答案:

答案 0 :(得分: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))

答案 1 :(得分:0)

下面的代码对我有用,但是对于较大的数字它很慢。

def compositions(k, n):
    def help_comp(k, n):
        C = set()
        if k == 1:
            for i in range(n):
                C.add((i+1,))
        else:
            for i in range(n):
                i=i+1
                for j in help_comp(k=k-1, n=n):
                    C.add((i,)+(j))
        return C
    C = set()
    if k == 1:
        C.add((n,))
    else:
        for i in range(n):
            i=i+1
            for j in help_comp(k=k-1, n=n):
                if sum(list((i,)+(j))) == n:
                    C.add((i,)+(j))
    return C