在Python中替换深层循环

Posted 2021-05-04

不幸的是，我经常在我的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指出它。

第1部分

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

第2部分

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表达了对它的赞赏......谁知道呢？