NumPy：获取 3D 数组行和的 argmax 的最快方法答案

【问题标题】：NumPy: Quickest way to get argmax of row sums of 3D arrayNumPy：获取 3D 数组行和的 argmax 的最快方法
【发布时间】：2019-06-27 00:11:13
【问题描述】：

假设我有一个 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 行，最快的方法是什么？

【问题讨论】：

标签： python numpy multidimensional-array

【解决方案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

【讨论】：