Recursivly looping through vectors with generators in Python

Question

The following python-code uses recursion to print all lists of nonnegative integers of length 3 whose sum is 2, it works as intended:

def rek(f,sum,n,vector=[]): #applies f to all Z^n_+ vectors of sum 'sum'
    if n==1:
        f(vector+[sum])
    else:        
        for i in range(sum+1):
            rek(f,sum-i,n-1,vector+[i])
                   
rek(print,2,3)

Output: 

 [0, 0, 2]  [0, 1, 1]  [0, 2, 0]  [1, 0, 1]  [1, 1, 0]  [2, 0, 0]

My question is if and how I could do this with generators instead? I would like to be able to write something like

for vector in vector_generator(2,3):
    print(vector)

to print the same vectors.

Answer 1

You should look into yield and yield from, here's how you'd do it:

def vector_generator(sum, n, vector=()):
    if n == 1:
        yield vector + (sum,)
    else:
        for i in range(sum + 1):
            yield from vector_generator(sum - i, n - 1, vector + (i,))

Note that I changed vector to a tuple, as giving a mutable object as a default argument is not good practice (editing the value will change the default argument).

Answer 2

You can convert your function to a generator in a simple fashion:

def vector_generator(sum,n,vector=None):
    vector = vector or []
    if n==1:
        yield vector+[sum]
    else:        
        for i in range(sum+1):
            yield from vector_generator(sum-i,n-1,vector+[i])


for vector in vector_generator(2,3):
    print(vector)

Answer 3

def gen(s,n):
    nums = [0 for _ in range(n)]
    while True:
        if sum(nums) == s:
            yield nums.copy()
        cidx = 0
        nums[cidx] += 1
        while nums[cidx] == s+1:
            nums[cidx] = 0
            cidx += 1
            if cidx == n:
                return
            nums[cidx] += 1
    
                   
for vector in gen(2,3):
    print(vector)

Not the most optimal solution but it should give correct results. What it does is basicly iterate over all the possible arrays of length n with numbers [0-s] and yields the ones that sum to s.