CODEWITHSUNDEEP
python
Pierpaolo
Pierpaolo

Reputation: 1741

alternative (faster) way to 3 nested for loop python

How can I make this function faster? (I call it a lot of times and it could result in some speed improvements)

def vectorr(I,  J,  K):
    vect = []
    for k in range(0,  K):
        for j in range(0, J):
            for i in range(0, I):
                vect.append([i, j, k])
    return vect

Upvotes: 5

Views: 6897

Answers (4)

greggo
greggo

Reputation: 3249

import numpy

def vectorr(I,J,K):
   val = numpy.indices( (I,J,K))
   val.shape = (3,-1)
   return val.transpose() # or val.transpose().tolist()

Upvotes: 1

Jimothy
Jimothy

Reputation: 9740

The other answers are more thorough and, in this specific case at least, better, but in general, if you're using Python 2, and for large values of I, J, or K, use xrange() instead of range(). xrange gives a generator-like object, instead of constructing a list, so you don't have to allocate memory for the entire list.

In Python 3, range works like Python 2's xrange.

Upvotes: 1

John Zwinck
John Zwinck

Reputation: 249394

You said you want it to be faster. Let's use NumPy!

import numpy as np

def vectorr(I, J, K):
    arr = np.empty((I*J*K, 3), int)
    arr[:,0] = np.tile(np.arange(I), J*K)
    arr[:,1] = np.tile(np.repeat(np.arange(J), I), K)
    arr[:,2] = np.repeat(np.arange(K), I*J)
    return arr

There may be even more elegant tweaks possible here, but that's a basic tiling that gives the same result (but as a 2D array rather than a list of lists). The code for this is all implemented in C, so it's very, very fast--this may be important if the input values may get somewhat large.

Upvotes: 2

Artsiom Rudzenka
Artsiom Rudzenka

Reputation: 29121

You can try to take a look at itertools.product

Equivalent to nested for-loops in a generator expression. For example, product(A, B) returns the same as ((x,y) for x in A for y in B).

The nested loops cycle like an odometer with the rightmost element advancing on every iteration. This pattern creates a lexicographic ordering so that if the input’s iterables are sorted, the product tuples are emitted in sorted order.

Also no need in 0 while calling range(0, I) and etc - use just range(I)

So in your case it can be:

import itertools

def vectorr(I,  J,  K):
    return itertools.product(range(K), range(J), range(I))

Upvotes: 8

Related Questions