Numpy：按另一个数组中的行对数组的行进行排序

Question

我有一个“邻居”的二维数组，我想根据另一个矩阵中的相应行（称为“半径”）对每一行重新排序。下面的代码有效，但它在一个 numpy 数组上使用了 for 循环，我知道这是不正确的方法。这种重新排序的正确 numpy / broadcast 解决方案是什么？

neighbors = np.array([[8,7,6], [3,2,1]])
radii = np.array([[0.4, 0.2, 0.1], [0.3, 0.9, 0.1]])

order = radii.argsort(axis=1)
for i in range(2):
    neighbors[i] = neighbors[i,order[i]]
print(neighbors)

# Result:
[[6 7 8]
 [1 3 2]]

Answer 1

在 NumPy 中你会这样写：

>>> neighbors[np.arange(2)[:, None], order]
array([[6, 7, 8],
       [1, 3, 2]])

（更一般地，您会将第一个索引写为np.arange(order.shape[0])[:, None]。）

这是因为np.arange(2)[:, None] 看起来像这样：

array([[0],
       [1]])

顺序如下：

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

对于花哨的索引，NumPy 对索引每个轴的数组进行配对。行索引[0] 与列索引[2, 1, 0] 配对，并按照此确定的顺序创建新行。同样对于[1] 和[2, 0, 1] 来确定第二行。