使用键数组根据第一个元素按顺序拉取元素

Question

我正在为以下问题寻找矢量化方法：

假设我有两个数组，一个在第一列中有一堆不连续的 id，在其余列中有一些数据，第二个数组建议我需要提取哪些数据线：

data_array = np.array([[101,4],[102,7],[201,2],[203,9],[403,12]])
key_array = np.array([101,403,201])

输出必须保持 key_array 给定的顺序，导致以下结果：

output_array = np.array([[101,4],[403,12],[201,2]])

我可以通过列表理解轻松做到这一点：

output_array = np.array([data_array[i==data_array[:,0]][0] for i in key_array])

但这不是矢量化解决方案。使用 numpy isin() 非常接近工作，但不保留给定的顺序：

data_array[np.isin(data_array[:,0],key_array)]

#[[101   4]
# [201   2] not the order given by the key_array!
# [403  12]]

我尝试通过使用 argsort() 来完成上述工作，但没有任何工作。任何帮助将不胜感激。

Answer 1

我们可以使用np.searchsorted -

s = data_array[:,0].argsort()
out = data_array[s[np.searchsorted(data_array[:,0],key_array,sorter=s)]]

如果data_array的第一列已经排序，则简化为单行-

out = data_array[np.searchsorted(data_array[:,0],key_array)]