Iterating over multiple indices with i > j ( > k) in a pythonic way

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

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:

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.

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)]

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)]

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

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)

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.

来源：https://stackoverflow.com/questions/41521725/iterating-over-multiple-indices-with-i-j-k-in-a-pythonic-way

标签

python

python-3.x

itertools