首页 > 解决方案 > NumPy 数组看起来很慢;难道我做错了什么?

问题描述

对于执行相同计算的这两个函数,为什么第二个函数需要大约八倍的时间来运行?

def random_walk_2(n):
    """Return coordinates after 'n' step random walk."""
    x, y = 0, 0
    for _ in range(n):
        (dx, dy) = random.choice([(0, 1), (0, -1), (1, 0), (-1, 0)])
        x += dx
        y += dy
    return (x, y)     # 6.48 seconds for 20,000

def random_walk_3(n):
    """Return coordinates after 'n' step random walk."""
    location = np.array([0, 0])
    for _ in range(n):
        move = np.array(random.choice([[0, 1], [0, -1], [1, 0], [-1, 0]]))
        location += move
    return location     # 54.8 seconds for 20,000

TIA,

标记

标签: numpy

解决方案


为了充分利用numpy,一次生成所有动作。在这里,我使用numpy's'的版本choice一次选择 1000 个动作。然后将它们相加:

In [138]: arr = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])
In [139]: moves = np.random.choice(4, 1000)
In [140]: np.sum(arr[moves,:], axis=0)
Out[140]: array([  9, -51])

另一个“运行”:

In [141]: moves = np.random.choice(4, 1000)
In [142]: np.sum(arr[moves,:], axis=0)
Out[142]: array([30, 34])

一个时机:

In [144]: timeit np.sum(arr[np.random.choice(4, 20000),:],0)
952 µs ± 190 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

推荐阅读