Swap pair of elements along an axis

user_6659511 0
Viewed 143

I have a 2d numpy array as such:

import numpy as np

a = np.arange(20).reshape((2,10))

# array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
#   [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

I want to swap pairs of elements in each row. The desired output looks like this:

# array([[ 9, 0, 2,  1,  4,  3,  6,  5,  8,  7],
#   [19, 10, 12, 11, 14, 13, 16, 15, 18, 17]])

I managed to find a solution in 1d:

a = np.arange(10)

# does the job for all pairs except the first
output = np.roll(np.flip(np.roll(a,-1).reshape((-1,2)),1).flatten(),2)

# first pair done manually
output[0] = a[-1]
output[1] = a[0]

Any ideas on a "numpy only" solution for the 2d case ?

python-3.x numpy
1 Answers

Owing to the first pair not exactly subscribing to the usual pair swap, we can do that separately. For the rest, it would relatively straight-forward with reshaping to split axes and flip axis. Hence, it would be -

In [42]: a # 2D input array
Out[42]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])

In [43]: b2 = a[:,1:-1].reshape(a.shape[0],-1,2)[...,::-1].reshape(a.shape[0],-1)

In [44]: np.hstack((a[:,[-1,0]],b2))
Out[44]: 
array([[ 9,  0,  2,  1,  4,  3,  6,  5,  8,  7],
       [19, 10, 12, 11, 14, 13, 16, 15, 18, 17]])

Alternatively, stack and then reshape+flip-axis -

In [50]: a1 = np.hstack((a[:,[0,-1]],a[:,1:-1]))

In [51]: a1.reshape(a.shape[0],-1,2)[...,::-1].reshape(a.shape[0],-1)
Out[51]: 
array([[ 9,  0,  2,  1,  4,  3,  6,  5,  8,  7],
       [19, 10, 12, 11, 14, 13, 16, 15, 18, 17]])
user_32938810
Related