alternative (faster) way to 3 nested for loop python

Question:

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
Asked By: Pierpaolo

||

Answers:

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))
Answered By: Artsiom Rudzenka

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.

Answered By: John Zwinck

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.

Answered By: Jimothy
import numpy

def vectorr(I,J,K):
   val = numpy.indices( (I,J,K))
   val.shape = (3,-1)
   return val.transpose() # or val.transpose().tolist()
Answered By: greggo
Categories: questions Tags:
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.