How to use numpy lexsort on a 3d array without changing the shape on the resulting array?

Question

我有一个像这样的 3d 数组：

arr = [
[[1, 2, 0], [1, 0, 1], [1, 0, 0]],
[[1, 2, 1], [1, 0, 2], [1, 2, 0]],
[[1, 2, 2], [2, 0, 0], [1, 0, 3]]
]

我想使用该长度为 3 的数组中的最后一项作为第一个排序列，然后是第二项，然后是第一项，对第二维数组中的内部数组进行排序。

所以结果应该是：

arr = [
[[1, 0, 0], [1, 2, 0], [1, 0, 1]],
[[1, 2, 0], [1, 2, 1], [1, 0, 2]],
[[2, 0, 0], [1, 2, 2], [1, 0, 3]]
]

我尝试这样做的方式如下：

arr = [
    [[1, 2, 0], [1, 0, 1], [1, 0, 0]],
    [[1, 2, 1], [1, 0, 2], [1, 2, 0]],
    [[1, 2, 2], [2, 0, 0], [1, 0, 3]]
]

sort = (arr[:, 2], arr[:, 1], arr[:, 0])
arr = arr[np.lexsort(sort)]

print(arr)

但是它把数组的形状从(3,3,3)改成了(3,3,3,3)，结果如下：

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

我试过了

sort = (arr[:, :, 2], arr[:, :, 1], arr[:, :, 0])
arr = arr[np.lexsort(sort)]

但它也将形状更改为 (3,3,3,3)

我错过了什么？

更新：

我设法让它像这样工作：

arr = np.array([
[[1, 2, 0], [1, 0, 1], [1, 0, 0]],
[[1, 2, 1], [1, 0, 2], [1, 2, 0]],
[[1, 2, 2], [2, 0, 0], [1, 0, 3]]
])

for i in range(len(arr)):
    item = arr[i]
    arr[i] = item[np.lexsort((item[:, 0], item[:, 1], item[:, 2]))]

print(arr)

但是这样的循环看起来不太好——我的实际数组非常非常大。大约 80Gb 的大内存。如果有人可以帮助我正确地调用 numpy 来获得正确的结果，我将非常感激！

Answer 1

花了整整一个小时让我头疼，在玩了大部分之后，我确信以下作品：

import numpy as np

arr = np.array([
    [[1, 2, 0], [1, 0, 1], [1, 0, 0]],
    [[1, 2, 1], [1, 0, 2], [1, 2, 0]],
    [[1, 2, 2], [2, 0, 0], [1, 0, 3]]
])

idx = np.lexsort(arr.T, axis=0)
out = arr[np.arange(arr.shape[0])[:, None], idx.T, :]

但我无法准确解释如何:)。如果排序不是简单的倒序（第 2 列，第 1 列，第 0 列），那么我需要多玩玩它......

出：

array([[[1, 0, 0],
        [1, 2, 0],
        [1, 0, 1]],

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

       [[2, 0, 0],
        [1, 2, 2],
        [1, 0, 3]]])

Answer 2

这实际上很简单：

arr = np.array([
[[1, 2, 0], [1, 0, 1], [1, 0, 0]],
[[1, 2, 1], [1, 0, 2], [1, 2, 0]],
[[1, 2, 2], [2, 0, 0], [1, 0, 3]]
])

for i in range(arr.shape[0]):
    arr[i] = np.moveaxis(arr[i], 1, -1)