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)]