如何对多维数组进行np.argsort排序

Question

如何使用np.argsort2和三维数组

例如：np.argsort1d数组工作：

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])

结果:

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]

我尝试在2d数组上应用np.argsort：

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])

结果出乎意料：

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]]]

因为我已经指定了axis=1，所以我希望np.argsort返回第二维度的idx。与wise一样，在3d数组上应用np.argsort会返回意外结果：

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])

结果:

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]]]]]

Answer 1

这个怎么样

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)

结果:

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]]