如何通过另一个2D数组对3D numpy数组的每一行进行排序？

Question

我有一个包含二维点的二维numpy数组：

np.random.seed(0)   
a = np.random.rand(3, 4, 2) # each value is a 2D point

我希望对每一行按照每个点的范数进行排序。

norms = np.linalg.norm(a, axis=2) # shape(3, 4)

indices = np.argsort(norms, axis=0) # indices of each sorted row

现在我想创建一个与a相同形状和值的数组，其中每一行的2D点按其范数排序。

我该如何实现呢？

我尝试了np.take和np.take_along_axis的不同变体，但都没有成功。

例如：

np.take(a, indices, axis=1) # shape (3,3,4,2)

这个示例将 a 按照 indices 的每一行分别进行了三次采样。

我希望只对 a 进行一次采样。每一行的 indices 列标识了应该从相应的行中进行采样的列。

- Elad Maimoni

---

当您修复问题后，我很乐意提名重新开放并发布答案。 - Mad Physicist

@MadPhysicist 我已经编辑了问题并提供了一个例子和澄清。希望现在清楚了。 - Elad Maimoni

@NinaKaprez，发布的答案是否不符合您的要求？如果不是，请提供一个样例输入和期望输出。 - Ehsan

带着轴怎么样？另外，你的argsort应该有axis=1吧？或者你是想对列进行排序而不是行？ - Mad Physicist

@MadPhysicist 这个怎么样？我没能使用它。我认为它不适合这个目的。我不明白问题出在哪里。 - Elad Maimoni

Answer 1

如果我理解你的意思正确，你想要这个：

norms = np.linalg.norm(a,axis=2) # shape(3,4)
indices = np.argsort(norms , axis=1)
np.take_along_axis(a, indices[:,:,None], axis=1)

您的示例输出结果：

[[[0.4236548  0.64589411]
  [0.60276338 0.54488318]
  [0.5488135  0.71518937]
  [0.43758721 0.891773  ]]

 [[0.07103606 0.0871293 ]
  [0.79172504 0.52889492]
  [0.96366276 0.38344152]
  [0.56804456 0.92559664]]

 [[0.0202184  0.83261985]
  [0.46147936 0.78052918]
  [0.77815675 0.87001215]
  [0.97861834 0.79915856]]]

- Ehsan

---

不，输出应该与变量'a'的形状相同。 - Elad Maimoni

@NinaKaprez 形状是相同的。 a 不是 (3,4,2) 的形状吗？如果不是，您需要更好地澄清。 - Ehsan

@Eshan，我的错。我想这就是我需要的。谢谢！你能解释一下indices[::None]是什么意思吗？ - Elad Maimoni

@NinaKaprez。它在“None”索引处添加了一个新的维度。None is np.newaxis。 - Mad Physicist

@NinaKaprez 不客气。不能比“疯狂物理学家”的评论说得更好了。 - Ehsan