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

Question

我是 Python 的新手，但仍然不能称自己为 Python 程序员。说到这里，如果我的问题没有任何意义，请多多包涵。

问题:

我有两个相同大小的 numpy 数组，例如A 和 B，其中 A.shape 等于 B.shape，它们都等于 (5,1000)，我想找到 A 中每一行的最大值以及 B 中相应元素的最大值。例如，如果在第四行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])