Function for compositions in Python

Question

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.

Answer 1

First, notice that the length of the solution grows very quickly with n, so the time and memory to compute this could blow in your face if you use large numbers.

Having said that, notice that getting all combinations of numbers and checking the sum is a bad idea because you are generating a lot of cases that by just checking the first numbers you know they won't work. In other words, you need to prune the search of arrays as quick as possible.

It is very funny to think about this kind of problems and better and faster implementations. Here you have one possibility:

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)

This is a generator, if you need the list, just do, for example, list(gen_combinations(3,4)).