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

时间:2015-07-30 20:39:46

标签: python arrays numpy

我是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)))

非常感谢任何指导我走正确方向的帮助。

3 个答案:

答案 0 :(得分:7)

以下是执行相同操作的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])

答案 1 :(得分:1)

作为bsxfun爱好者,很高兴看到有人试图将相同的功能复制到其他编程语言中。现在,bsxfun基本上是broadcasting机制,它也存在于NumPy中。在NumPy中,通过使用np.newaxis或简单None创建单例维度来实现。

回到上下文中的问题,可以实现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])

答案 2 :(得分:0)

print np.max(A[i])这将在i矩阵的numpy行中给出最高值。