查找numpy数组中每行的最大值以及相同大小的另一个数组中的相应元素

Question

我是Python的新手,仍然不能称自己为Python程序员.说到这一点,如果我的问题没有任何意义,请耐心等待.

题:

我有两个相同大小的numpy数组,例如A和B,其中A.shape等于B.shape并且它们都相等(5,1000),我想找到A中每行的最大值和相应的元素例如,如果在A的第四行中,最大元素索引是104,那么我想在数组B中找到第四行的第104个元素,而对于其余行则相同.

我知道我可以通过循环遍历行来实现它,但我想知道是否有更优雅的方式.例如,如果我在MATLAB中这样做,我会写下面的代码:

B(bsxfun(@eq,A,max(A,[],2)))

任何指导我通过正确方向的帮助将非常感激.

Answer 1

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