CODEWITHSUNDEEP
pythonlistfor-loopnested
NNsr
NNsr

Reputation: 1063

Neat way of writing nested for loop in python

(Update) I need to find stationary distribution of a Markov Chain with 4.5 million states. It means I need to solve a linear system with 4.5 million equations. Each state is a vector of size 6. I am trying to store each state in a list. The following in part of my effort in creating all admissible states.

I am trying to loop through a big set of numbers and create a set of vectors. Here is a simplified version of my code:

mylist=[]
for i in range(1,4):
    for j in range(1,4-i):
        for k in range(0,5-i-j):
            Temp=[i,j,k]
            mylist.extend(Temp)
            print(mylist)
            mylist=[]
            Temp=[]

which will give me:

[1, 1, 0]
[1, 1, 1]
[1, 1, 2]
[1, 2, 0]
[1, 2, 1]
[2, 1, 0]
[2, 1, 1]

my question is: is there a neater, nicer, more efficient way of doing this in Python?

Thank you

Upvotes: 0

Views: 958

Answers (3)

Ravichandra
Ravichandra

Reputation: 2342

from itertools import product
mylist =[[i,j,k] for i,j,k in product(range(1,4),range(1,4),range(2))]

Upvotes: 0

Anand S Kumar
Anand S Kumar

Reputation: 90889

If you are looking for a one line code that would create the same vector , you can use list comprehension in python.

Example -

myList = [[i,j,k] for i in range(1,4) for j in range(1,4-i) for k in range(0,5-i-j)]
myList
>> [[1, 1, 0], [1, 1, 1], [1, 1, 2], [1, 2, 0], [1, 2, 1], [2, 1, 0], [2, 1, 1]]

Though I do not think this is in anyway neater or more efficient.

Though after some testing using timeit , we can see that list comprehension may be a little bit faster -

In [1]: def foo1():
   ...:     l = []
   ...:     for i in range(100):
   ...:                 for j in range(100):
   ...:                         l.append([i,j])
   ...:     return l

In [3]: def foo2():
   ...:     return [[i,j] for i in range(100) for j in range(100)]
   ...: 

In [4]: %timeit foo1()
100 loops, best of 3: 3.08 ms per loop

In [5]: %timeit foo2()
100 loops, best of 3: 2.16 ms per loop

In [6]: %timeit foo2()
100 loops, best of 3: 2.18 ms per loop

In [7]: %timeit foo1()
100 loops, best of 3: 3.11 ms per loop

Upvotes: 3

Rafael T
Rafael T

Reputation: 15679

I completely don't get what you are trying to archive with your lists. You could get the same output with this:

for i in range(1,4):
    for j in range(1,4-i):
        for k in range(0,5-i-j):
            print([i,j,k])

Upvotes: 1

Related Questions