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

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

1 个答案: