不幸的是，我经常在我的Python代码中使用while循环导致程序显着减速。
以下是while循环（shape = (1000,1000,3)）的示例：
i = 0
j = 0
while i < arr.shape[0]:
while j < arr.shape[1]:
if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
arr[i,j,:] = 1
else:
arr[i,j,:] = 0
j = j + 1
j = 0
i = i + 1
-->
def f(x):
return 1 if x <= 5and x > 0 else 0
f = np.vectorize(f)
arr= f(arr)
编辑：另一个while循环
i = 0
j = 0
while i < arr1.shape[0]:
while j < arr1.shape[1]:
if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
arr1[i, j, :] = arr1[i, j, :]
else:
arr1[i, j, :] = arr2[i, j, :]
j = j + 1
j = 0
i = i + 1
有没有办法加快速度？我不知道怎么做。
编辑我很急，我之前的回答是不正确的。感谢@gboffi指出它。
def original(arr):
i = 0
j = 0
while i < arr.shape[0]:
while j < arr.shape[1]:
if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
arr[i,j,:] = 1
else:
arr[i,j,:] = 0
j = j + 1
j = 0
i = i + 1
return arr
def vectorized(arr):
mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
arr[mask] = 1
arr[~mask] = 0
return arr
def vectorized2(arr):
"""Works only if assigning 0 and 1s"""
mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
mask = np.dstack([mask] * arr.shape[2])
return mask.astype(np.float32)
基准
以下使用下一个基准测试中的arr。我没有时间对此进行大量测试，因此我建议您运行下面的矢量化版本并使用np.allclose来比较原始代码的结果，以便您可以验证arr的各种情况。
In [101]: %timeit v1 = vectorized(arr.copy())
1000 loops, best of 3: 312 µs per loop
In [102]: %timeit v2 = original(arr.copy())
100 loops, best of 3: 10.4 ms per loop
In [103]: np.allclose(v1, v2)
Out[103]: True
In [108]: %timeit v3 = vectorized2(arr.copy())
10000 loops, best of 3: **83.3 µs** per loop
In [110]: v3 = vectorized2(arr.copy())
In [111]: np.allclose(v1, v3)
Out[111]: True
def vectorized(arr1, arr2):
mask = np.all(arr1 == 0, axis=2)
mask = np.dstack([mask] * arr1.shape[2])
return np.where(mask, arr1, arr2)
def original(arr1, arr2):
i = 0
j = 0
while i < arr1.shape[0]:
while j < arr1.shape[1]:
if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
arr1[i, j, :] = arr1[i, j, :]
else:
arr1[i, j, :] = arr2[i, j, :]
j = j + 1
j = 0
i = i + 1
return arr1
第二个循环的基准
# Prepare data
m = 100
n = 100
d = 3
np.random.seed(0)
arr = np.random.randint(0, 11, size=(m, n, d))
true_mask = np.random.randint(0, 2, size=(m, n, 1), dtype=np.bool)
true_mask = np.dstack([true_mask] * d)
arr[true_mask] = 0
arr1 = arr.copy()
arr2 = -1 * np.ones_like(arr1)
In [84]: v1 = vectorized(arr1, arr2)
In [85]: v2 = original(arr1, arr2)
In [86]: np.allclose(v1, v2)
Out[86]: True
In [87]: %timeit v1 = vectorized(arr1, arr2)
1000 loops, best of 3: 284 µs per loop
In [88]: %timeit v2 = original(arr1, arr2)
100 loops, best of 3: 12.6 ms per loop
让我们生成一些测试数据
In [61]: np.random.seed(0) ; a = np.random.randint(10, size=(3,4,2))-3
In [62]: a
Out[62]:
array([[[ 2, -3],
[ 0, 0],
[ 4, 6],
[ 0, 2]],
[[-1, 1],
[ 4, 3],
[ 5, 5],
[-2, 3]],
[[ 4, 4],
[ 5, -2],
[ 2, 6],
[ 5, 6]]])
根据问题中的示例，第一个元素在[1 ... 5]中的每一行变为一行，所有其他行变为零行。
让我们生成一个反映OP正在执行的测试的掩码（仅涉及上面显示中每行的第一个元素）
In [63]: t = np.logical_and(0<a[:,:,0],a[:,:,0]<=5)
并使用此掩码来处理行...这里我们必须添加冒号运算符来完成寻址
In [64]: a[ t,:] = 1
In [65]: a[~t,:] = 0
让我们显示结果
In [66]: a
Out[66]:
array([[[1, 1],
[0, 0],
[1, 1],
[0, 0]],
[[0, 0],
[1, 1],
[1, 1],
[0, 0]],
[[1, 1],
[1, 1],
[1, 1],
[1, 1]]])
如果我正确解释OP示例，这就是他们想要的。
我不得不说这个答案显然与我的相似
In [67]: np.random.seed(0) ; a = np.random.randint(10, size=(3,4,2))-3
In [68]: t = (0 < a) & (a <= 5)
In [69]: a[t] = 1
In [70]: a[~t]= 0
In [71]: a
Out[71]:
array([[[1, 0],
[0, 0],
[1, 0],
[0, 1]],
[[0, 1],
[1, 1],
[1, 1],
[0, 1]],
[[1, 1],
[1, 0],
[1, 0],
[1, 0]]])
给出了不同的结果，OP表达了对它的赞赏......谁知道呢？