numpy argmax 和 sort 函数的时间比较

Question

我试图在 numpy 中找到二维矩阵的列总和最大的列。例如：

令 A = [[1, 2, 3], [0, 1, 4], [0, 0, 1]]

每一列的总和是[1, 3, 8]。因此，第 3 列具有最大列和。

在尝试 numpy.argmax 和 numpy.sort 函数来完成此任务时，我期望 argmax 比理想情况下的 sort 更快，但它们导致相同的运行时间。

a = np.random.rand(7000, 8000)
start_time = time.time()
for i in range(1000):
    np.sort(np.sum(a, axis = 0))
print(time.time() - start_time)

上面的代码运行了 33.29 秒，而下面的代码也运行了 34.33 秒。

a = np.random.rand(7000, 8000)
start_time = time.time()
for i in range(1000):
    np.argmax(np.sum(a, axis=0))
print(time.time() - start_time)

你能告诉我这背后的潜在原因吗？是否与我解决问题的方式有关？

Answer 1

使用 timeit 模块计时。

>>> from timeit import Timer
>>> import numpy as np
>>> a = np.random.random((7000,8000))
>>> loops = 3
>>> timer = Timer("np.sum(a, axis=0)", "from __main__ import a, np")
>>> timer.timeit(loops) / loops
0.10155341827648574
>>> timer = Timer("np.argmax(a)", "from __main__ import a, np")
>>> timer.timeit(loops) / loops
0.11956859843814982
>>> timer = Timer("np.sort(a)", "from __main__ import a, np")
>>> timer.timeit(loops) / loops
3.5973468146321466
>>> timer = Timer("np.sort(np.sum(a, axis=0))", "from __main__ import a, np")
>>> timer.timeit(loops) / loops
0.09826639265653132
>>> timer = Timer("np.argmax(np.sum(a, axis=0))", "from __main__ import a, np")
>>> timer.timeit(loops) / loops
0.09442937388683958
>>>

numpy argmax 和 sort 函数的时间比较

Time comparison of numpy argmax and sort functions

python

numpy

python-3.x

numpy-ndarray