NumPy 数组看起来很慢；难道我做错了什么？

Question

对于执行相同计算的这两个函数，为什么第二个函数的运行时间大约是其八倍？

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,
标记

Answer 1

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

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)