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.
CodePudding user response:
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).
CodePudding user response:
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)
CodePudding user response:
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
.