Найти максимум каждой строки в массиве numpy и соответствующий элемент в другом массиве того же размера
Я новичок в Python и до сих пор не могу назвать себя программистом Python. Говоря об этом, пожалуйста, потерпите меня, если мой вопрос не имеет никакого смысла.
вопрос:
У меня есть два массива numpy одинакового размера, например A и B, где A. shape равен B. shape и они оба равны (5,1000), и я хочу найти максимальное значение каждой строки В A и соответствующий элемент этого в B. Например, если в четвертой строке a максимальный индекс элемента равен 104, то я хотел бы найти 104-й элемент четвертой строки в массиве B и то же самое для остальных строк.
Я знаю, что могу сделать это, зацикливаясь на строках, но мне было интересно, есть ли более элегантный способ сделать это. Например, если бы я сделал это в MATLAB, я бы написал следующий код:
B(bsxfun(@eq,A,max(A,[],2)))
любая помощь, которая проведет меня через правильное направление, будет очень признательна.
3 ответов
здесь numpy идиома для того, чтобы делать то же самое:
b[np.arange(len(a)), np.argmax(a, axis=1)]
например:
>>> a = np.array([
[1, 2, 0],
[2, 1, 0],
[0, 1, 2]
])
>>> b = np.array([
[1, 2, 3],
[1, 2, 3],
[1, 2, 3]
])
>>> b[np.arange(len(a)), np.argmax(a, axis=1)]
array([2, 1, 3])
будучи bsxfun любовник, здорово видеть, как люди пытаются воспроизвести ту же функциональность на других языках программирования. Теперь,bsxfun в основном broadcasting механизм, который существует и в NumPy. В NumPy это достигается путем создания одноэлементных измерений с помощью np.newaxis или просто None.
вернемся к вопросу в контексте, an equivalent решение на основе широковещания может быть реализовано, как показано в примере запуска -
In [128]: A
Out[128]:
array([[40, 63, 67, 65, 19],
[85, 55, 66, 92, 88],
[50, 1, 23, 6, 59],
[67, 55, 46, 78, 3]])
In [129]: B
Out[129]:
array([[78, 63, 45, 34, 81],
[ 5, 38, 28, 61, 66],
[ 3, 65, 16, 25, 32],
[72, 1, 31, 75, 6]])
In [130]: B[A == A.max(axis=1)[:,None]]
Out[130]: array([45, 61, 32, 75])