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.
Iterating over a list can also be achieved using a while loop. The block of code inside the loop executes until the condition is true. A loop variable can be used as an index to access each element.
Use the map() Function to Iterate Over Two Lists in Python The map() function accepts iterable data structures such as lists and strings as input along with a function. After applying the specified function, it maps the iterable data elements and returns a tuple after mapping the values.
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)
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With