14

i need to iterate over a tuple of indices. all indices must be in the range [0, N) with the condition i > j. The toy example I present here deals with only two indices; I will need to extend that to three (with i > j > k) or more.

The basic version is this:

N = 5
for i in range(N):
    for j in range(i):
        print(i, j)

and it works just fine; the output is

1 0
2 0
2 1
3 0
3 1
3 2
4 0
4 1
4 2
4 3

I don't want to have one more indentation level for every additional index, therefore I prefer this version:

for i, j in ((i, j) for i in range(N) for j in range(i)):
    print(i, j)

this works perfectly well, does what it should and gets rid of the extra indentation level.

I was hoping to be able to have something more elegant (for two indices that is not all that relevant, but for three or more it becomes more relevant). What I came up with so far is this:

from itertools import combinations

for j, i in combinations(range(N), 2):
    print(i, j)

This iterates over the same pair of indices just fine. The only thing that is different is the order in which the pairs appear:

1 0
2 0
3 0
4 0
2 1
3 1
4 1
3 2
4 2
4 3

As the order of what I am doing with these indices is relevant, I therefore cannot use this.

Is there an elegant, short, pythonic way to iterate over these indices in the same order that the very first example produces? Keep in mind that N will be large, so sorting is not something I would want to do.

hiro protagonist
  • 44,693
  • 14
  • 86
  • 111

5 Answers5

17

You could solve this generally as follows:

def indices(N, length=1):
    """Generate [length]-tuples of indices.

    Each tuple t = (i, j, ..., [x]) satisfies the conditions 
    len(t) == length, 0 <= i < N  and i > j > ... > [x].

    Arguments:
      N (int): The limit of the first index in each tuple.
      length (int, optional): The length of each tuple (defaults to 1).

    Yields:
      tuple: The next tuple of indices.

    """
    if length == 1:
       for x in range(N):
           yield (x,)
    else:
       for x in range(1, N):
            for t in indices(x, length - 1):
                yield (x,) + t

In use:

>>> list(indices(5, 2))
[(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2), (4, 0), (4, 1), (4, 2), (4, 3)]
>>> list(indices(5, 3))
[(2, 1, 0), (3, 1, 0), (3, 2, 0), (3, 2, 1), (4, 1, 0), (4, 2, 0), (4, 2, 1), (4, 3, 0), (4, 3, 1), (4, 3, 2)]
jonrsharpe
  • 115,751
  • 26
  • 228
  • 437
  • Maybe [`yield from`](https://docs.python.org/3/whatsnew/3.3.html#pep-380) (supported from Python 3.3) could make the code more compact, like `yield from range(N)` – lazzzis Jan 07 '17 at 14:46
  • @lazzzis I thought about that, but it needs to do something to each value before yielding it – jonrsharpe Jan 07 '17 at 14:50
6

You can use product from itertools if you don't mind the inefficiency of throwing out most of the generated tuples. (The inefficiency gets worse as the repeat parameter increases.)

>>> from itertools import product
>>> for p in ((i,j) for (i,j) in product(range(5), repeat=2) if i > j):
...   print p
...
(1, 0)
(2, 0)
(2, 1)
(3, 0)
(3, 1)
(3, 2)
(4, 0)
(4, 1)
(4, 2)
(4, 3)
>>> for p in ((i,j,k) for (i,j,k) in product(range(5), repeat=3) if i > j > k):
...   print p
...
(2, 1, 0)
(3, 1, 0)
(3, 2, 0)
(3, 2, 1)
(4, 1, 0)
(4, 2, 0)
(4, 2, 1)
(4, 3, 0)
(4, 3, 1)
(4, 3, 2)

Update: Instead of tuple unpacking, using indexing for the filter. This allows the code to be written a little more compactly. Only my_filter needs to be changed for tuples of varying sizes.

from itertools import product, ifilter
def my_filter(p):
    return p[0] > p[1] > p[2]

for p in ifilter(my_filter, product(...)):
    print p
chepner
  • 497,756
  • 71
  • 530
  • 681
  • that looks very clean indeed. thanks! will have to time how bad the inefficiency gets... – hiro protagonist Jan 07 '17 at 14:40
  • You can alleviate some of the inefficiency with `product(*[range(5-x) for x in range(3)])` (where `5` and `3` are modified as desired, of course), but that detracts from the elegance without really solving the issue entirely. – chepner Jan 07 '17 at 15:30
6

Here's an approach with itertools.combinations to have a generic number of levels -

map(tuple,(N-1-np.array(list(combinations(range(N),M))))[::-1])

Or a bit twisted one with same method -

map(tuple,np.array(list(combinations(range(N-1,-1,-1),M)))[::-1])

, where N : number of elements and M : number of levels.

Sample run -

In [446]: N = 5
     ...: for i in range(N):
     ...:     for j in range(i):
     ...:         for k in range(j):  # Three levels here
     ...:             print(i, j, k)
     ...:             
(2, 1, 0)
(3, 1, 0)
(3, 2, 0)
(3, 2, 1)
(4, 1, 0)
(4, 2, 0)
(4, 2, 1)
(4, 3, 0)
(4, 3, 1)
(4, 3, 2)

In [447]: N = 5; M = 3

In [448]: map(tuple,(N-1-np.array(list(combinations(range(N),M))))[::-1])
Out[448]: 
[(2, 1, 0),
 (3, 1, 0),
 (3, 2, 0),
 (3, 2, 1),
 (4, 1, 0),
 (4, 2, 0),
 (4, 2, 1),
 (4, 3, 0),
 (4, 3, 1),
 (4, 3, 2)]
Divakar
  • 218,885
  • 19
  • 262
  • 358
  • aha! rearranging the output from `combinations`. looks interesting! – hiro protagonist Jan 07 '17 at 14:41
  • @hiroprotagonist Well its not meant to be a generator as I am using it as a NumPy array. So, to get the final output, we need to wrap it up with tuple : `tuple(map(tuple,(N-1-np.array(list(combinations(range(N),M))))[::-1]))`. That answers your question? – Divakar Jan 07 '17 at 14:52
  • @hiroprotagonist Or if you don't mind an array of indices as the final output : `N-1-np.array(list(combinations(range(N),M)))[::-1]`. – Divakar Jan 07 '17 at 14:54
  • the drawback is that you create a list and have all the indices in memory; it's no longer a 'pure' generator, right? unfortunately `islice` can not handle negative steps... – hiro protagonist Jan 07 '17 at 14:56
  • @hiroprotagonist No, it's not a generator once we have wrapped it with `np.array()` there and thus creates all output indices in memory. Would that be a problem? – Divakar Jan 07 '17 at 15:00
  • the only 'issue' is the memory (more a conceputal issue than a practical one...). because then i could just define: `def ijk(n, depth): lst = reversed(list(combinations(range(n-1, -1, -1), depth))); yield from lst` and get the indices from there. (same idea; no `numpy`). – hiro protagonist Jan 07 '17 at 15:04
  • @hiroprotagonist Yeah that should work, but still creates those indices in memory it seems. – Divakar Jan 07 '17 at 15:20
2

This is an approach based on the observation that it is easier to generate the negatives of the indices in the (reverse of) the desired order It is similar to the approach of @Divakar and like that has the drawback of requiring the list to be created in memory:

def decreasingTuples(N,k):
    for t in reversed(list(itertools.combinations(range(1-N,1),k))):
        yield tuple(-i for i in t)

>>> for t in decreasingTuples(4,2): print(t)

(1, 0)
(2, 0)
(2, 1)
(3, 0)
(3, 1)
(3, 2)
>>> for t in decreasingTuples(4,3): print(t)

(2, 1, 0)
(3, 1, 0)
(3, 2, 0)
(3, 2, 1)
John Coleman
  • 51,337
  • 7
  • 54
  • 119
  • yup, that is the same as the one in my comment in [Divakar's answer](http://stackoverflow.com/a/41522251/4954037) (the one with `def ijk(n, depth)`), right? still like it! – hiro protagonist Jan 07 '17 at 16:05
  • ah, no, now i see the difference! – hiro protagonist Jan 07 '17 at 16:10
  • @hiroprotagonist Slightly different, since I am using a range which loops through negative values -- if it was exactly the same I wouldn't have posted. jon sharpe's seems best. – John Coleman Jan 07 '17 at 16:10
-1

a somewhat 'hacky' attempt using eval (just adding this for completeness. there are nicer answers here!).

the idea is to construct a string like

'((a, b, c) for a in range(5) for b in range(a) for c in range(b))'

and return the eval of that:

def ijk_eval(n, depth):
    '''
    construct a string representation of the genexpr and return eval of it...
    '''

    var = string.ascii_lowercase
    assert len(var) >= depth > 1  # returns int and not tuple if depth=1

    for_str = ('for {} in range({}) '.format(var[0], n) +
               ' '.join('for {} in range({})'.format(nxt, cur)
                        for cur, nxt in zip(var[:depth-1], var[1:depth])))
    return eval('(({}) {})'.format(', '.join(var[:depth]), for_str))

can be used this way and produces the right results.

for i, j in ijk_eval(n=5, depth=2):
    print(i, j)

the construction is not very nice - but the result is: it is a regular genexpr and just as efficient as those are.

hiro protagonist
  • 44,693
  • 14
  • 86
  • 111