Вытягивание элементов по порядку на основе первого элемента с использованием массива ключей

Question

Я ищу векторизованный подход для следующей задачи:

Предположим, у меня есть два массива, один с набором несмежных идентификаторов в первом столбце и некоторыми данными в оставшихся столбцах, а второй массив, предлагающий, какие данные мне нужно извлечь:

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