Python: how to avoid loop?

One way in which you should be able to significantly speed up your code is to avoid duplicate work:

tmp = np.where(M == i) 

As this gives you a list of all locations in M where the value is equal to i, it must be searching through the entire matrix. So for each element in l, you are searching through the full matrix.

Instead of doing that, try indexing your matrix as a first step:

matrix_index = {} for i in len(M): for j in len(M[i]): if M[i][j] not in matrix_index: matrix_index[M[i][j]] = [(i,j)] else: matrix_index[M[i][j]].append((i,j)) 

Then for each value in l, instead of doing a costly search through the full matrix, you can just get it straight from your matrix index.

_{Note: I haven't with numpy very much, so I may have gotten the specific syntax incorrect. There may also be a more idiomatic way of doing this in numpy.}

If both l and M are not large matrices like the following:

 In: l0 = [5, 3, 8, 12, 34, 1, 12] In: M0 = [[12, 34, 5, 8, 7], In: [ 0, 24, 12, 3, 1]] In: l = np.asarray(l) In: M = np.asarray(M) 

You can try this:

 In: np.where(l[None, None, :] == M[:, :, None]) Out: (array([0, 0, 0, 0, 0, 1, 1, 1, 1]), <- i array([0, 0, 1, 2, 3, 2, 2, 3, 4]), <- j array([3, 6, 4, 0, 2, 3, 6, 1, 5])) <- k 

The rows should be the i, j, k, respectively and read the column to get every (i, j, k) you need. For example, the 1st column [0, 0, 3] means M[0, 0] = l[3], and the 2nd column [0, 0, 6] says M[0, 0] = l[6], and vice versa. I think these are what you want.

However, the numpy trick can not be extended to very large matrices, such as 2M elements in l or 2500x2500 elements in M. They need quite a lot memory and very very long time to compute... if they are lucky not to crash for out of memory. :)

One solution that does not use the word for is

c = np.apply_along_axis(lambda row: np.random.choice(np.argwhere(row).ravel()), 1, M.ravel()[np.newaxis, :] == l[:, np.newaxis]) indI, indJ = c // M.shape[1], c % M.shape[1] 

Note that while that solves the problem, M.ravel()[np.newaxis, :] == l[:, np.newaxis] will quickly produce MemoryErrors. A more pragmatic approach would be to get the indices of interest through something like

s = np.argwhere(M.ravel()[np.newaxis, :] == l[:, np.newaxis]) 

and then do the random choice post-processing by hand. This, however, probably does not yield any significant performance improvements over your search.

What makes it slow, though, is that you search through the entire matrix in every step of your loop; by pre-sorting the matrix (at a certain cost) gives you a straightforward way of making each individual search much faster:

In [312]: %paste def direct_search(M, l): indI = [] indJ = [] for i in l: tmp = np.where(M == i) rd = np.random.randint(len(tmp[0])) # Note the fix here indI.append(tmp[0][rd]) indJ.append(tmp[1][rd]) return indI, indJ def using_presorted(M, l): a = np.argsort(M.ravel()) M_sorted = M.ravel()[a] def find_indices(i): s = np.searchsorted(M_sorted, i) j = 0 while M_sorted[s + j] == i: yield a[s + j] j += 1 indices = [list(find_indices(i)) for i in l] c = np.array([np.random.choice(i) for i in indices]) return c // M.shape[1], c % M.shape[1] ## -- End pasted text -- In [313]: M = np.random.randint(0, 1000000, (1000, 1000)) In [314]: l = np.random.choice(M.ravel(), 1000) In [315]: %timeit direct_search(M, l) 1 loop, best of 3: 4.76 s per loop In [316]: %timeit using_presorted(M, l) 1 loop, best of 3: 208 ms per loop In [317]: indI, indJ = using_presorted(M, l) # Let us check that it actually works In [318]: np.all(M[indI, indJ] == l) Out[318]: True