I have an array n x m x k. The task is to sort 1st axis by selected 2nd axis element. Currently i do it like in example below. The input array is 3x3x3, axis 1 is sorted by 1st element of 2nd axis. Is there more efficient way using nditer?
import numpy as np
a = np.random.rand(27).reshape(3,3,3)
print(a)
for i in np.arange(np.shape(a)[0]):
idx = np.argsort(a[i,:,1])
a[i]=a[i,idx]
print(a)
Approach #1
Get the argsort indices selecting all elements along the first axis in one-go and then use NumPy's advanced-indexing to index along the second axis to get the re-arranged output -
idx = np.argsort(a[:,:,1])
a_out = a[np.arange(a.shape[0])[:,None], idx]
Approach #2
If we are doing in-situ edits in the input array (writing back the results to input), we can run a loop after pre-computing those argsort indices, like so -
idx = np.argsort(a[:,:,1])
for i,indx in enumerate(idx):
a[i] = a[i,indx]
# Original method
In [130]: np.random.seed(0)
...: a = np.random.rand(100,100,100)
In [131]: %%timeit
...: for i in np.arange(np.shape(a)[0]):
...: idx = np.argsort(a[i,:,1])
...: a[i]=a[i,idx]
1000 loops, best of 3: 1.63 ms per loop
# Approach #1
In [132]: np.random.seed(0)
...: a = np.random.rand(100,100,100)
In [133]: %%timeit
...: idx = np.argsort(a[:,:,1])
...: a_out = a[np.arange(a.shape[0])[:,None], idx]
1000 loops, best of 3: 1.6 ms per loop
# Approach #2
In [134]: np.random.seed(0)
...: a = np.random.rand(100,100,100)
In [135]: %%timeit
...: idx = np.argsort(a[:,:,1])
...: for i,indx in enumerate(idx):
...: a[i] = a[i,indx]
1000 loops, best of 3: 1.24 ms per loop