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