numpy: 按指定值重新排序数组

Question

我有一个矩阵：

A = [ [1,2],
      [3,4],
      [5,6] ]

还有一个数值向量：

V = [4,6,2]

我想根据第二列的值，使用向量V中的值来重新排列矩阵A。结果应该是：

A = [ [3,4],
      [5,6], 
      [1,2] ] # 2nd columns' values have the same order as V

该怎么做呢？

Answer 1

@JoeKington的numpy解决方案非常聪明，但它依赖于A[:,1]的顺序是排好序的。这里有一个适用于一般情况的修正方法：

import numpy as np

np.random.seed(1)
N=5
A = np.arange(2*N).reshape((-1,2))+100
np.random.shuffle(A)
print(A)

如果A看起来像这样：

[[104 105]
 [102 103]
 [108 109]
 [100 101]
 [106 107]]

而V

V = A[:,1].copy()
np.random.shuffle(V)
print(V)

看起来像这样：

[105 109 107 101 103]

那么我们可以使用Joe的方法：

vals, order = np.unique(np.hstack((A[:,1],V)), return_inverse=True)

但要同时保存A[:,1]和V的顺序：

a_order = order[:V.size]
v_order = order[-V.size:]

然后在用v_order重新排序之前，先对A进行排序（通过形成A[np.argsort(a_order)]）：

print A[np.argsort(a_order)][v_order]

[[104 105]
 [108 109]
 [106 107]
 [100 101]
 [102 103]]

(A[np.argsort(a_order)]是根据它的第二列排序后的A。)

---

注意，np.unique总是以排序的顺序返回数组。文档保证当return_inverse=True时，返回的索引是唯一数组的索引，这些索引可以重建原始数组。也就是说，如果你这样调用np.unique：

uniq_arr, indices = np.unique(arr, return_inverse=True)

你可以保证

unique_arr[indices] = arr

因为你可以依赖这个关系，Joe的方法并不依赖于一个简单的实现细节——unique总是会这样工作。（这可是名言——考虑到np.unique1d返回的输出参数的顺序发生了什么……不过这不重要 :))

Answer 2

首先，我们需要找到数组 A 第二列中，与 V 的顺序匹配的值的索引。在这个例子中，这些索引是 [1,2,0]。一旦我们得到了这些索引，就可以使用 numpy 的“花式”索引来完成剩下的操作。

所以，你可以这样做：

import numpy as np
A = np.arange(6).reshape((3,2)) + 1
V = [4,6,2]
column = A[:,1].tolist()
order = [column.index(item) for item in V]
print A[order,:]

如果你想完全避免使用 Python 列表，那么你可以尝试下面的做法。这种方法有点黑科技，可能还有更好的方法...

我们可以利用 numpy.unique 来实现这个功能... 我这里做的事情依赖于一个特定的实现细节（unique 似乎是从数组的末尾开始的），这个细节随时可能会改变... 这就是为什么这是一种不太优雅的黑科技。

import numpy as np
A = np.arange(6).reshape((3,2)) + 1
V = np.array([4,6,2])
vals, order = np.unique(np.hstack((A[:,1],V)), return_inverse=True)
order = order[-V.size:]
print A[order,:]