Swap multiple indices in 2d array

Question

Problem

For the following array:

import numpy as np

arr = np.array([[i for i in range(10)] for j in range(5)])

# arr example
array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
       [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])

For each row in arr, I'm attempting to swap n (in this case 2) indices according to some 2d array, such as.

swap = np.random.choice(arr.shape[1], [arr.shape[0], 2], replace=False)

# swap example
array([[8, 1],
       [5, 0],
       [7, 2],
       [9, 4],
       [3, 6]])

Question

I tried arr[:, swap] = arr[:, swap[:, ::-1]], but this performs every swap for each row, rather than only swapping indices row by row. The behaviour I am trying to achieve is given below. Is this possible without iterating over swap?

for idx, s in enumerate(swap):
   arr[idx, s] = arr[idx, s[::-1]]

# new arr with indices swapped
array([[0, 8, 2, 3, 4, 5, 6, 7, 1, 9],
       [5, 1, 2, 3, 4, 0, 6, 7, 8, 9],
       [0, 1, 7, 3, 4, 5, 6, 2, 8, 9],
       [0, 1, 2, 3, 9, 5, 6, 7, 8, 4],
       [0, 1, 2, 6, 4, 5, 3, 7, 8, 9]])

iz_ · Accepted Answer · 2020-01-27 18:58:56Z

You can to use a "helper" array to index arr. The helper coerces arr into the correct shape.

import numpy as np

arr = np.array([[i for i in range(10)] for j in range(5)])
swap = np.array([[8, 1], [5, 0], [7, 2], [9, 4], [3, 6]])

helper = np.arange(arr.shape[0])[:, None]
# helper is
# array([[0],
#        [1],
#        [2],
#        [3],
#        [4]])
# arr[helper] is
# array([[[0, 8, 2, 3, 4, 5, 6, 7, 1, 9]],
#        [[5, 1, 2, 3, 4, 0, 6, 7, 8, 9]],
#        [[0, 1, 7, 3, 4, 5, 6, 2, 8, 9]],
#        [[0, 1, 2, 3, 9, 5, 6, 7, 8, 4]],
#        [[0, 1, 2, 6, 4, 5, 3, 7, 8, 9]]])

arr[helper, swap] = arr[helper, swap[:, ::-1]]

# arr is
# array([[0, 8, 2, 3, 4, 5, 6, 7, 1, 9],
#        [5, 1, 2, 3, 4, 0, 6, 7, 8, 9],
#        [0, 1, 7, 3, 4, 5, 6, 2, 8, 9],
#        [0, 1, 2, 3, 9, 5, 6, 7, 8, 4],
#        [0, 1, 2, 6, 4, 5, 3, 7, 8, 9]])

Collectives™ on Stack Overflow

Swap multiple indices in 2d array

Problem

Question

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

Problem

Question

1 Answer 1

Comments

Related