How to rotate a square numpy array with different times efficiently by `np.rot90`?

Question

I have a 2d numpy array, for example:

a = np.array([
        [0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])

and another 1d array:

I = np.array([0, 2, 3, 1, 0, 2, 0, 1])

I want to rotate a by np.rot90 function like following:

b = np.zeros((len(I), 3, 3))
for i, k in enumerate(I):
    b[i] = np.rot90(a, k=k)

Can I do it more efficiently without the floop?

Answer 1

Approach #1

Generate a 3D array of all possible 4 rotations and simply index into it with I and thus have a vectorized solution -

P = np.empty((4,) + a.shape, dtype=a.dtype)
P[0] = a            # For np.rot90(a, k=0)
P[1] = a.T[::-1]    # For np.rot90(a, k=1)
P[2] = a[::-1,::-1] # For np.rot90(a, k=2)
P[3] = a.T[:,::-1]  # For np.rot90(a, k=3)
out = P[I]

Approach #2

Another way to create P would be with -

P = np.array([np.rot90(a, k=i) for i in range(4)])

and as with the previous method simply index into P with I for final output.

Runtime test

Approaches -

def org_app(a, I):
    m,n = a.shape
    b = np.zeros((len(I), m, n), dtype=a.dtype)
    for i, k in enumerate(I):
        b[i] = np.rot90(a, k=k)
    return b

def app1(a, I):
    P = np.empty((4,) + a.shape, dtype=a.dtype)
    P[0] = a  
    P[1] = a.T[::-1]
    P[2] = a[::-1,::-1]
    P[3] = a.T[:,::-1]
    return P[I]

def app2(a, I):
    P = np.array([np.rot90(a, k=i) for i in range(4)])
    return P[I]

Timings -

In [54]: a = np.random.randint(0,9,(10,10))

In [55]: I = np.random.randint(0,4,(10000))

In [56]: %timeit org_app(a, I)
10 loops, best of 3: 51 ms per loop

In [57]: %timeit app1(a, I)
1000 loops, best of 3: 469 µs per loop

In [58]: %timeit app2(a, I)
1000 loops, best of 3: 549 µs per loop

100x+ speedup!

Answer 2

我能想到的（仍然没有向量化的）一种更有效的方法是在一行中使用列表理解：

np.array([np.rot90(a, k=i) for i in I])