通过二维索引对三维 numpy 数组进行排序答案

【问题标题】：Sorting a three-dimensional numpy array by a two-dimensional index通过二维索引对三维 numpy 数组进行排序
【发布时间】：2019-11-04 09:16:07
【问题描述】：

我在这方面遇到了真正的麻烦。我有一个三维 numpy 数组，我想通过二维索引数组对其进行重新排序。实际上，数组将以编程方式确定，并且三维数组可能是二维或四维的，但为了简单起见，如果两个数组都是二维的，则以下是所需的结果：

ph = np.array([[1,2,3], [3,2,1]])
ph_idx = np.array([[0,1,2], [2,1,0]])
for sub_dim_n, sub_dim_ph_idx in enumerate(ph_idx):
    ph[sub_dim_n] = ph[sub_dim_n][sub_dim_ph_idx]

这使得 ph 数组变成：

array([[1, 2, 3],
       [1, 2, 3]])

这是我想要的。如果情况相同，谁能帮忙，但是我有一个三维数组（psh）而不是 ph，例如：

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

希望这很清楚，如果不是，请询问。提前致谢！

【问题讨论】：

在您的三维案例中，ph_idx 的示例是什么？
@fuglede - 会是（对此处的格式表示歉意）np.array([[[1,2,3]],[[1,2,3]]])
格式化是最好的。这次不是[3, 2, 1]？您能否提供一个ph 形状为(2, 2, 3) 的示例？
所以，我没有提出这个问题，因为我不想让事情变得过于复杂，但 ph 的 dim1 的长度可能大于 1（如 2、2、3 形状你建议）。如果是这样，我正在尝试应用二维数组。在类似 numpy 伪代码的东西中，这将是 ph[ph_idx[:, None, :]]

标签： python arrays numpy

【解决方案1】：

如果您想要以 ph.shape 形状的数组结尾，您可以简单地 np.squeeze ph_ixs 使形状匹配，并使用它来索引 ph：

print(ph)
[[[1 2 3]]
 [[3 2 1]]]

print(ph_idx)
[[0 1 2]
 [2 1 0]]

np.take_along_axis(np.squeeze(ph), ph_idx, axis=-1)

array([[1, 2, 3],
       [1, 2, 3]])

【讨论】：

ph 是他们示例中的 3 维数组，而不是 ph_idx。也就是说，np.take_along_axis 的一些变体可能是他们所追求的。
这对我来说是一个新功能！至少在 1.10 上不可用。
谢谢 - 我同意 - take_along_axis 看起来像在这里使用的函数。
@Chrisper：如果您的所有示例都遵循您概述的内容，那么您甚至都不需要np.squeeze；普通的np.take_along_axis 完成了这项工作。
@yatu：如果你使用cmets到OP的ph_idx，你甚至可以把np.squeeze放到ph上。

【解决方案2】：

所以，线索已经在这里有用的 cmets 中，但为了完整起见，它就像使用 np.take_along_axis 和 2d 数组的广播版本一样简单：

psh = np.array(
    [[[1,2,3]], 
     [[3,2,1]]]
)
ph_idx = np.array(
    [[0,1,2], 
     [2,1,0]]
)
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

如果 3d 数组的 dim1 包含多个元素，这也具有以下优点：

psh = np.array(
    [[[1,2,3],[4,5,6],[7,8,9]], 
     [[3,2,1],[6,5,4],[9,8,7]]]
)
ph_idx = np.array([[0,1,2], [2,1,0]])
np.take_along_axis(psh, ph_idx[:, None, :], axis=2)

给了

array([[[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]],

       [[1, 2, 3],
        [4, 5, 6],
        [7, 8, 9]]])

【讨论】：