Numpy 阵列上的两个 for 循环

Question

早上好。 假设我有一个使用 numpy 创建的二维数组（称为 MAT(x,y)）。 在这个数组上，我必须执行一些操作。 如何重写以下 2 个 for 循环，例如使用np.nditer()或其他使用 numpy 方法的东西？ 谢谢你。

    for i in range(x): 
        for j in range(y): 

            if i == 0: MAT[i][j] = j   
            elif j == 0: MAT[i][j] = i

Answer 1

您可以像这样简单地设置第一行和第一列

mat[:,0] = np.arange(0, mat.shape[0])
mat[0,:] = np.arange(0, mat.shape[1])

示例结果

array([[0.        , 1.        , 2.        , 3.        , 4.        ],
       [1.        , 0.30487009, 0.97179858, 0.08143348, 0.99363866],
       [2.        , 0.69357714, 0.98421733, 0.42032313, 0.81041628]])

Answer 2

由于您只是将值 0,1,... 分配给第一行和第一列，因此不需要内部带有 if 条件的双循环。

只需将值分配给第一行：

MAT[0] = np.arange(len(MAT[0])

并到第一列：

MAT[:,0] = np.arange(len(MAT[:,0]))

Answer 3

您基本上可以将np.arange()分配给输入的适当切片。

您可以对输入的 2D 特性进行硬编码（ foo2() ），也可以使用动态定义的切片（ foon() ）对任意维度执行此操作：

import numpy as np


def foo(arr):
    ii, jj = arr.shape
    for i in range(ii): 
        for j in range(jj): 
            if i == 0:
                arr[i, j] = j   
            elif j == 0:
                arr[i, j] = i
    return arr


def foo2(arr):
    ii, jj = arr.shape
    arr[:, 0] = np.arange(ii)
    arr[0, :] = np.arange(jj)
    return arr


def foon(arr):
    for i, d in enumerate(arr.shape):
        slicing = tuple(slice(None) if j == i else 0 for j in range(arr.ndim))
        arr[slicing] = np.arange(d)
    return arr


arr = np.zeros((3, 4))

print(foo(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

print(foo2(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

print(foon(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

请注意，双重切片（例如MAT[i][j] ）在工作时不如使用元组切片（例如MAT[i, j] ）高效。

---

最后，循环的嵌套在您的代码中基本上没有使用，您可以通过分开的两个循环来重写它（这样效率更高）：

def fool(arr):
    ii, jj = arr.shape
    for i in range(ii):
        arr[i, 0] = i
    for j in range(jj):
        arr[0, j] = j
    return arr

这很有趣，因为如果我们使用 Numba 加速代码：

fool_nb = nb.jit(fool)
fool_nb.__name__ = 'fool_nb'

这导致最快的方法：

funcs = foo, foo2, foon, fool, fool_nb


shape = 300, 400
for func in funcs:
    arr = np.zeros(shape)
    print(func.__name__)
    %timeit func(arr)
    print()

# foo
# 100 loops, best of 3: 6.53 ms per loop

# foo2
# 100000 loops, best of 3: 4.28 µs per loop

# foon
# 100000 loops, best of 3: 6.99 µs per loop

# fool
# 10000 loops, best of 3: 89.8 µs per loop

# fool_nb
# 1000000 loops, best of 3: 1.01 µs per loop

Answer 4

您不需要循环，只需分配给切片

MAT[:, 0] = np.arange(x)
MAT[0, :] = np.arange(y)