CODEWITHSUNDEEP
pythonpython-3.xnumpyoptimization
Gekke Boer
Gekke Boer

Reputation: 33

Efficient NumPy rows rotation over variable distances

Given a 2D M x N NumPy array and a list of rotation distances, I want to rotate all M rows over the distances in the list. This is what I currently have:

import numpy as np

M = 6
N = 8
dists = [2,0,2,1,4,2] # for example
matrix = np.random.randint(0,2,(M,N))

for i in range(M):
    matrix[i] = np.roll(matrix[i], -dists[i])

The last two lines are actually part of an inner loop that gets executed hundreds of thousands of times and it is bottlenecking my performance as measured by cProfile. Is it possible to, for instance, avoid the for-loop and to do this more efficiently?

Upvotes: 3

Views: 92

Answers (1)

Divakar
Divakar

Reputation: 221744

We can simulate the rolling behaviour with modulus operation after adding dists with a range(0...N) array to give us column indices for each row from where elements are to be picked and shuffled in the same row. We can vectorize this process across all rows with the help of broadcasting. Thus, we would have an implementation like so -

M,N = matrix.shape # Store matrix shape

# Get column indices for all elems for a rolled version with modulus operation
col_idx = np.mod(np.arange(N) + dists[:,None],N)

# Index into matrix with ranged row indices and col indices to get final o/p
out = matrix[np.arange(M)[:,None],col_idx]

Upvotes: 1

Related Questions