向量化 for 循环 Python

Question

我有一个嵌套的 for 循环：

import numpy as np

ccounter = np.zeros(shape=(120, 200))
lat_idx = np.random.randint(120, size=4800)
lon_idx = np.random.randint(200, size=(4800, 4800))
for j in range(4800):
    for i in range(4800):
        ccounter[lat_idx[i], lon_idx[i, j]] +=1

这显然很慢。是否可以避免 for 循环并将其实现为例如矩阵运算？

Answer 1

这是一个矢量化方法 np.bincount -

# Get matrix extents for output
k = lon_idx.max()+1 # 200 for given sample
n = lat_idx.max()+1 # 120 for given sample

# Get linear index equivalent
lidx = lat_idx[:,None]*k+lon_idx

# Use those indices as bins for binned count. Reshape for final o/p
out = np.bincount(lidx.ravel(),minlength=n*k).reshape(n,k)

为了进一步提高大型数组的性能，我们可以利用 numexpr 得到 lidx -

import numexpr as ne

lidx = ne.evaluate('lat_idx2D*k+lon_idx',{'lat_idx2D':lat_idx[:,None]})

向量化 for 循环 Python

Vectorize for loop Python

python

for-loop

numpy

vectorization