How to np.argsort multi-dimensional array

user_2611189 0
Viewed 53

How to np.argsort 2 and 3-dimensional array?

For example: np.argsort 1d array works:

x = np.random.randint(0,3, (10,))
print("before", x.shape)
print(x)
print()
idx = np.argsort(x, axis=0)
print("idx", idx.shape)
print(idx)
print()
print("after", x[idx].shape)
print(x[idx])

Result:

before (10,)
[1 2 2 0 0 1 1 1 0 0]

idx (10,)
[3 4 8 9 0 5 6 7 1 2]

after (10,)
[0 0 0 0 1 1 1 1 2 2]

My attempt on applying np.argsort on 2d array:

x = np.random.randint(0,3, (2, 10,))
print("before", x.shape)
print(x)
print()
idx = np.argsort(x, axis=1) # this returns wrong results, see the next code snippet for results
print("idx", idx.shape)
print(idx)
print()
print("after", x[:,idx].shape)
print(x[:,idx])

The result is unexpected:

before (2, 10)
[[0 2 0 1 1 1 2 1 1 2]
 [2 0 2 0 2 1 0 0 2 0]]

idx (2, 10)
[[0 2 3 4 5 7 8 1 6 9] # wrong idx sequence
 [1 3 6 7 9 5 0 2 4 8]] # I believe this should be: [6 0 7 1 8 5 2 3 9 4]

after (2, 2, 10) # expected shape (2, 10)
[[[0 0 1 1 1 1 1 2 2 2]
  [2 1 2 1 2 1 0 0 1 1]]

 [[2 2 0 2 1 0 2 0 0 0]
  [0 0 0 0 0 1 2 2 2 2]]]

Since I have specified axis=1, I am expecting np.argsort to return the idx for the 2nd dimension. Like wise, applying np.argsort on 3d array returns unexpected results:

x = np.random.randint(0,3, (2, 1, 10,))
print("before", x.shape)
print(x)
print()
idx = np.argsort(x, axis=2)
print("idx", idx.shape)
print(idx)
print()
print("after", x[:,:,idx].shape)
print(x[:,:,idx])

Result:

before (2, 1, 10)
[[[1 2 1 1 0 1 0 2 2 0]]

 [[0 0 2 1 1 1 2 2 2 0]]]

idx (2, 1, 10)
[[[4 6 9 0 2 3 5 1 7 8]] # wrong idx sequence

 [[0 1 9 3 4 5 2 6 7 8]]] # wrong idx sequence

after (2, 1, 2, 1, 10) # expected shape (2, 1, 10)
[[[[[0 0 0 1 1 1 1 2 2 2]]

   [[1 2 0 1 0 1 1 0 2 2]]]]

 [[[[1 2 0 0 2 1 1 0 2 2]]

   [[0 0 0 1 1 1 2 2 2 2]]]]]
numpy python
1 Answers

How about this?

x = np.random.randint(0,3, (2, 10,))
print("before", x.shape)
print(x)
print()
idx = np.argsort(x, axis=1)
print("idx", idx.shape)
print(idx)
print()
after = np.take_along_axis(x, idx, axis=1)
print("after", after.shape)
print(after)

Result:

before (2, 10)
[[2 2 0 1 1 1 2 0 1 1]
 [2 2 0 1 0 1 1 2 2 0]]

idx (2, 10)
[[2 7 3 4 5 8 9 0 1 6]
 [2 4 9 3 5 6 0 1 7 8]]

after (2, 10)
[[0 0 1 1 1 1 1 2 2 2]
 [0 0 0 1 1 1 2 2 2 2]]
user_148449360
Related