python - NumPy:获取 3D 数组行和的 argmax 的最快方法

Question

假设我有一个 3 维数组，尺寸为 10x10000x5。将此数组解释为 10 个“子数组”，每个子数组有 10000 行和 5 列，我想要做的是，对于每一行:

(1) 计算 10 个子数组中每个子数组中行的总和。

(2) 确定哪个子数组的总和最高。

下面显示了一个示例。我执行上述操作，但仅针对前两行，其中“firstrow”是每个子数组的第一行的总和，“secondrow”是每个子数组的第二行的总和。然后，我使用 np.argmax() 查找产生最高总和的子数组。但我想对所有 10000 行执行此操作，而不仅仅是前两行。

import numpy as np
np.random.seed(777)
A = np.random.randn(10,10000,5)

first = [None]*10
second = [None]*10
for i in range(10):
    firstrow[i] = A[i].sum(axis=1)[0]
    secondrow[i] = A[i].sum(axis=1)[1]

np.argmax(np.array(firstrow)) # Sub-array 9 yields the highest sum
np.argmax(np.array(secondrow)) # Sub-array 8 yields the highest sum
#...

对于所有 10000 行，最快的方法是什么？

Answer 1

你可以这样做:

result = A.sum(2).argmax(0)

在您的示例中进行测试:

import numpy as np

np.random.seed(777)
A = np.random.randn(10, 10000, 5)

result = A.sum(2).argmax(0)

# Check against loop
first = [None] * 10
second = [None] * 10
for i in range(10):
    first[i] = A[i].sum(axis=1)[0]
    second[i] = A[i].sum(axis=1)[1]

print(result[0], np.argmax(np.array(first)))
# 9 9
print(result[1], np.argmax(np.array(second)))
# 8 8