python - 选择 numpy 列中的最大元素并获取它的行

Question

我正在尝试执行以下操作，给定一个 numpy 数组“A”，我需要获取每一列的最大元素(以绝对值表示)并获取该元素所在的行(重复值无关紧要) .

import numpy as np

A = np.array(
    [[1, 1, 1],
     [2, 1, 3],
     [3, 1, 6]], dtype=float)

到目前为止，我天真的解决方案如下:

for i in range(len(A)):
    max_elm_col = A[0,i]
    max_elm_row = 0
    for index, el in enumerate(A[0:, i]):
        if abs(el) > max_elm_col:
            max_elm_col = el
            max_elm_row = index

    print(max_elm_col, max_elm_row)

哪些输出:

3.0 2
1.0 0
6.0 2 

虽然这确实解决了我的问题，但看起来更 (num)Pythonic。所以我试着像这样修改它:

for i in range(len(A)):
    # cool way
    max_elm_col = max(A[0:,i], key = abs)
    print(max_elm_col)

哪些输出:

3.0
1.0
6.0

检查 docs 后，同样可以用 np.amax 完成，但它不允许我使用 key = abs。

for i in range(len(A)):
    max_elm_col = np.amax(A[0:,i], axis=0)
    print(max_elm_col)

由于我不熟悉 numpy，所以我希望获得有关以更 numpyish 的方式执行此任务的任何提示。

Answer 1

一种方法是计算堆叠数组中的 abs 值，例如:

代码:

maxs = np.vstack([A, np.abs(A)]).argmax(axis=0) % A.shape[0]

测试代码:

import numpy as np

A = np.array(
    [[1,  1, 1],
     [2, -2, 3],
     [2,  1, 3],
     [3,  1, 6]], dtype=float)


maxs = np.vstack([A, np.abs(A)]).argmax(axis=0) % A.shape[0]
print(maxs)
print(A[maxs, np.arange(A.shape[1])])

结果:

[3 1 3]
[ 3. -2.  6.]