Python：使用Numpy更快地生成2D数组

Question

我想学习使用Numpy在Python中生成2D数组的更快方法

我当前的代码使用生成器：

h,w= img.shape[:2]
points= [[x-w/2,y-h/2,img[y,x]] for y in range(h) for x in range(w) if img[y,x]!=0]

问题很简单：我们如何使用Numpy加快这一速度

注：img为2D numpy.array，其大小为10 x 10。上面的代码本身并不慢，但是我需要重复336次，这需要时间

非常感谢,

Answer 1

如果我正确理解了你的问题，你可以使用纯numpy解决方案

设置示例数据

import numpy as np
np.random.seed(10)
img = np.random.randint(0,10,(10,10))
h,w= img.shape
points= [[x-w/2,y-h/2,img[y,x]] for y in range(h) for x in range(w) if img[y,x]!=0]
points = np.array(points)

与努比

np_points = np.vstack([np.where(img != 0), img[img != 0].ravel()]).T[:,[1,0,2]] - [w/2 , h/2, 0]
assert np.array_equal(np_points,points), 'solutions are not equal'

---

对这两种解决方案进行基准测试没有太大区别。纯numpy解决方案可能会受益于一次性对所有336幅图像进行矢量化计算

%%timeit
for i in range(336):
  np_points = np.vstack([np.where(img != 0), img[img != 0].ravel()]).T[:,[1,0,2]] - [w/2 , h/2, 0]

输出：

100 loops, best of 3: 7.71 ms per loop

%%timeit
for i in range(336):
  points= [[x-w/2,y-h/2,img[y,x]] for y in range(h) for x in range(w) if img[y,x]!=0]

输出：

10 loops, best of 3: 25.6 ms per loop