我正在寻找解决以下问题的矢量化方法:
假设我有两个数组，一个数组的第一列中有一堆不连续的 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() 来实现上述工作，但没有成功。任何帮助将不胜感激。
最佳答案
我们可以使用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)]
关于python - 使用键数组根据第一个元素按顺序拉取元素，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/56976236/