How to turn this function into a generator?

Scott David Daniels scott.daniels at acm.org
Fri May 9 14:32:26 EDT 2003


Anton Vredegoor wrote:

> Please have a look at the following function:
 > ... I'd like to replace "res.append(a[:N])" with "yield a[:N]".the

Here's a 'mechanical' conversion with changed code ###'ed out
and followed by replacement.
> def genPartitions(m):
>     ### res,a = [],[0]*m
       a = [0]*m
>     def recPartition(m,B,N):
>             if m == 0:
>                 ### res.append(a[:N])
                   yield a[:N]
>             else:
>                 for i in range(1,min(B,m)+1):
>                     a[N]=i
>                     ### recPartition(m-i,i,N+1)
                       for partition in recPartition(m-i,i,N+1):
                           yield partition
>     ### recPartition(m,m,0)
>     ### return res
       for partition in recPartition(m,m,0):
           yield partition


And here's how I actually prefer:

def distribute(arr, left, pos):
     """generate all partitions beginning with arr[:pos] with totals of
     left for the arr[pos:] portion"""
     if left:
         for size in range(1, min(arr[pos-1], left)+1):
             arr[pos] = size
             for partition in distribute(arr, left-size, pos+1):
                     yield partition
     else:
	yield arr[:pos]  # nothing left to distribute

def generate(total):
     """generate all distinct partitions of total into pieces"""
     result = [0] * total
     for size in range(total+1):
	result[0] = size
         for partition in distribute(result, total-size, 1):
             yield partition

for t in generate(7):
     print t

-Scott David Daniels
Scott.Daniels at Acm.Org





More information about the Python-list mailing list