Skip to content

Instantly share code, notes, and snippets.

@mckoss
Created November 3, 2009 17:16
Show Gist options
  • Save mckoss/225235 to your computer and use it in GitHub Desktop.
Save mckoss/225235 to your computer and use it in GitHub Desktop.
Some nice Tuple-Iterators for Python
class tuples(object):
"""
Iterator that generates all tuples in lexicographical order.
For example:
for t in tuples((2,3)):
print t
produces:
(0, 0)
(0, 1)
(0, 2)
(1, 0)
(1, 1)
(1, 2)
"""
def __init__(self, aLimits):
self.aLimits = aLimits
self.size = len(aLimits)
self.aNext = [0 for i in range(self.size)]
self.fFirst = True
def __iter__(self):
return self
def next(self):
if self.fFirst:
self.fFirst = False
return tuple(self.aNext)
i = self.size - 1
while True:
if i < 0:
raise StopIteration
self.aNext[i] += 1
if self.aNext[i] >= self.aLimits[i]:
self.aNext[i] = 0
i -= 1
continue
return tuple(self.aNext)
def diagonal(n):
"""
Return all tuples of order n starting from (0,0,...0). The order
is diagonalized.
Example:
for t in diagonal(3):
print t
produces:
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(1, 0, 0)
(0, 0, 2)
(0, 1, 1)
(0, 2, 0)
(1, 0, 1)
(1, 1, 0)
(2, 0, 0)
(0, 0, 3)
(0, 1, 2)
...
"""
s = 0
while True:
fs = fixed_sum(n, s)
for t in fs:
yield t
s += 1
def fixed_sum(n, s):
"""
Return all tuples of order n where the sum of the elements equals s
"""
n = int(n)
s = int(s)
if n < 1:
raise Exception("fixed_sum order n out of bounds: %d" % n)
if s < 0:
raise Exception("fixed_sum sum s is out of bounds: %d" %s)
if n == 1:
yield (s,)
return
for i in range(0,s+1):
fs = fixed_sum(n-1, s-i)
for t in fs:
r = [i]
r.extend(t)
yield tuple(r)
if __name__ == '__main__':
print "tuples"
for t in tuples((2,3)):
print t
print "fixed_sum"
fs = fixed_sum(2,7)
for t in fs:
print t
print "diagonals"
c = 35
for t in diagonal(3):
print t
c -= 1
if c == 0:
break;
@azatsman
Copy link

Nice code -- thank you.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment