还有另一种按“列”值对 3D numpy 数组进行排序的方法吗？

Posted 2023-02-23

【中文标题】还有另一种按“列”值对 3D numpy 数组进行排序的方法吗？

【英文标题】：is there another way for sorting 3D numpy array by 'column' values?

【发布时间】：2020-08-10 11:30:51

【问题描述】：

我编写了一个 Python 代码，根据第一列的值对一组四个 3x3 表进行排序。有没有更简单的方法，那就是用更少的代码，也许更高效？ 这是我的代码：

import numpy as np 

np.random.seed(4)

a = np.random.randint(10, size=(4, 3, 3))
ind = a[:,:,0].argsort()
ind = np.stack(a.shape[2]*[ind], axis=1)
b = np.take_along_axis(a.transpose(0, 2, 1), ind, axis=2).transpose(0, 2, 1)

print(a)
print("----------------")
print(b)

[[[7 5 1]
  [8 7 8]
  [2 9 7]]

 [[7 7 9]
  [8 4 2]
  [6 4 3]]

 [[0 7 5]
  [5 9 6]
  [6 8 2]]

 [[5 8 1]
  [2 7 0]
  [8 3 1]]]
----------------
[[[2 9 7]
  [7 5 1]
  [8 7 8]]

 [[6 4 3]
  [7 7 9]
  [8 4 2]]

 [[0 7 5]
  [5 9 6]
  [6 8 2]]

 [[2 7 0]
  [5 8 1]
  [8 3 1]]]

【参考方案1】：

你有一个很好的解决方案。这是一个更短（可能更快）的：

b = np.einsum('iijk->ijk', a[:,a[:,:,0].argsort()])

einsum 基本上可以通过索引实现您想要实现的目标。它采用i-th 的a[:,a[:,:,0].argsort() 元素的i-th 元素。

b:

[[[2 9 7]
  [7 5 1]
  [8 7 8]]

 [[6 4 3]
  [7 7 9]
  [8 4 2]]

 [[0 7 5]
  [5 9 6]
  [6 8 2]]

 [[2 7 0]
  [5 8 1]
  [8 3 1]]]

【讨论】：

非常感谢 Ehsan，优雅的解决方案，而且速度更快。 ``` 我的：每个循环 22.8 µs ± 682 ns（平均值 ± 7 次运行的标准偏差，每个循环 10000 个循环）你的：每个循环 6.13 µs ± 123 ns（平均值 ± 7 次运行的标准偏差，每个循环 100000 个循环) ```

@Deliosiret 很高兴它有帮助。如果它解决了问题，请继续接受答案，以便其他人也觉得它有帮助。谢谢。