Выберите значения из массива numpy на основе другого массива

Question

Скажите, что у меня есть два массива a и b:

a = np.array([[1,2,3], [4,5,6], [7,8,9]])
b = np.array([[3,1,0], [1,2,3], [3,0,2]])

Я хочу выбрать из каждой строки в a элемент, соответствующий наибольшему значению (в пределах строки ) в b, т.е. я хочу вывод [1, 6, 7].

Каким было бы быстрое решение этой проблемы в pandas / NumPy, и было бы это быстрее, чем использование for-l oop в обычном python? Это кажется очень простым, но я не нашел хорошего решения. Я новичок в pandas / NumPy, но я думаю, что должно быть простое решение для этого?

Answer 1

Вы можете использовать np.argmax с axis=1, чтобы получить индекс каждого наибольшего значения в строках b,

, а затем использовать расширенную индексацию, чтобы получить элементы, которые вы хочу от a, вот так:

import numpy as np

a = np.array([[1,2,3], [4,5,6], [7,8,9]])
b = np.array([[3,1,0], [1,2,3], [3,0,2]])

b_largest_idx = np.argmax(b, axis=1)

print(a[range(a.shape[0]),b_largest_idx])

Вывод:

[1 6 7]