Friendsnippets

1. Generate all partitions of an integer

   view plain - save this

   Iterative algorithms to generate all partitions of a given integer. Only tested with Python 2.5

     1 import itertools
     2 
     3 class FixedLengthPartitionIterator:
     4     """Iterate over partitions of integer `N` into *exactly* `K`
     5     positive integers.
     6 
     7     Each returned partition is a list of positive integers in
     8     descending order, such that their sum is `N`.
     9 
    10     Arguments `min_` and `max_` bound the values in each partition.
    11 
    12     Examples::
    13       >>> list(FixedLengthPartitionIterator(3,1))
    14       [(3,)]
    15       >>> list(FixedLengthPartitionIterator(3,2))
    16       [(2, 1)]
    17       >>> list(FixedLengthPartitionIterator(3,3))
    18       [(1, 1, 1)]
    19       >>> list(FixedLengthPartitionIterator(6,2))
    20       [(5, 1), (4, 2), (3, 3)]
    21       >>> list(FixedLengthPartitionIterator(8,3))
    22       [(6, 1, 1), (5, 2, 1), (4, 3, 1), (4, 2, 2), (3, 3, 2)]
    23       >>> list(FixedLengthPartitionIterator(8,4,2))
    24       [(2, 2, 2, 2)]
    25       >>> list(FixedLengthPartitionIterator(8,3,2))
    26       [(4, 2, 2), (3, 3, 2)]
    27       >>> list(FixedLengthPartitionIterator(8,3,2,3))
    28       [(3, 3, 2)]      
    29     """
    30     def __init__(self, N, K, min_=1, max_=None):
    31         # `max_` really defaults to N
    32         if max_ is None:
    33             max_ = N
    34 
    35         # rule out trivial cases
    36         if (K*min_ > N) or (K*max_ < N):
    37             self.done = True
    38             return
    39 
    40         self._N = N  #: integer to be partitioned
    41         self._K = K  #: maximum number of nonzero parts
    42         self._k = 0  #: current number of nonzero parts
    43         self._min = min_  #: minimum value of each part
    44         self._max = max_  #: maximum value of each part
    45         self._p = [0] * K #: current partition
    46         self.done = False #: when `True`, enumeration is over
    47 
    48     def __iter__(self):
    49         return self
    50 
    51     def next(self):
    52         if self.done:
    53             raise StopIteration
    54         if self._N == self._K * self._min:
    55             self.done = True
    56             return tuple([self._min]*self._K)
    57         else:
    58             while self._k <= self._K:
    59                 i = self._k - 1
    60                 while i > 0:
    61                     if (self._p[i]+1 <= self._p[i-1]-1) \
    62                            and (self._p[i-1]-1 >= self._min) \
    63                            and (self._p[i]+1 <= self._max):
    64                         self._p[i-1] -= 1
    65                         self._p[i] += 1
    66                         return tuple(self._p)
    67                     else:
    68                         i -= 1
    69                 # only change the first `k` parts
    70                 self._k += 1
    71                 head = (self._N - self._k - self._min*self._K + self._min + 1)
    72                 if (head < self._min+1) or (head > self._max) \
    73                        or (self._k > self._K):
    74                     continue
    75                 # advance to next partition:
    76                 # [N-2*(k-1)-(K-k), 2, ..., 2, (k-1 times) 1, ..., 1 (K-k times)]
    77                 self._p = [head] \
    78                           + ([self._min + 1] * (self._k - 1)) \
    79                           + ([self._min] * (self._K - self._k))
    80                 return tuple(self._p)
    81             raise StopIteration
    82 
    83 
    84 def PartitionIterator(N, K, min_=1, max_=None):
    85     """Iterate over partitions of integer `N` into *at most* `K`
    86     positive integers.
    87 
    88     Each returned partition is a list of positive integers in
    89     descending order, such that their sum is `N`.
    90 
    91     Optional arguments `min_` and `max_` bound the values in each
    92     partition.
    93 
    94     Examples::
    95       >>> list(PartitionIterator(2,1))
    96       [(2,)]
    97       >>> list(PartitionIterator(3,3))
    98       [(3,), (2, 1), (1, 1, 1)]
    99       >>> list(PartitionIterator(8,3,2))
   100       [(8,), (6, 2), (5, 3), (4, 4), (4, 2, 2), (3, 3, 2)]
   101       >>> list(PartitionIterator(8,3,2,3))
   102       [(3, 3, 2)]      
   103     """
   104     return itertools.chain(*[FixedLengthPartitionIterator(N,k,min_,max_)
   105                              for k in xrange(1,K+1)])
   

   Posted by rmurri to python algorithms mathematics partitions combinatorics ... saved by 2 persons ... 0 comments ... 1 year

Pages : 1