python - numpy argmax 和 sort 函数的时间比较

Question

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

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

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

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

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)

您能否让我知道这背后的潜在原因？这与我如何解决问题有关吗？

标签： pythonpython-3.xnumpynumpy-ndarray

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
>>>