跳转至

NumPy迭代

NumPy 提供了多种方法来迭代 ndarray 数组。可以对数组的元素进行逐个处理，或对多维数组的特定维度进行迭代。以下是一些常用的迭代方法及其详细示例：

直接迭代

对 ndarray 进行直接迭代会遍历数组的第一维。

import numpy as np

# 一维数组的迭代
arr_1d = np.array([1, 2, 3, 4, 5])
for element in arr_1d:
    print(element)
# 输出:
# 1
# 2
# 3
# 4
# 5

# 二维数组的迭代
arr_2d = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
for row in arr_2d:
    print(row)
# 输出:
# [1 2 3]
# [4 5 6]
# [7 8 9]

nditer 迭代器

nditer 是 NumPy 提供的多维数组迭代器，适用于多种复杂迭代需求。

基本用法

arr = np.array([[1, 2, 3], [4, 5, 6]])
for element in np.nditer(arr):
    print(element)
# 输出:
# 1
# 2
# 3
# 4
# 5
# 6

迭代时修改数组

for element in np.nditer(arr, op_flags=['readwrite']):
    element[...] = element * 2
print(arr)
# 输出:
# [[ 2  4  6]
#  [ 8 10 12]]

控制迭代顺序

默认情况下，nditer 按行优先（C 风格）顺序迭代。可以通过指定 order 参数控制迭代顺序。

arr = np.array([[1, 2, 3], [4, 5, 6]])

# 按列优先（Fortran 风格）顺序迭代
for element in np.nditer(arr, order='F'):
    print(element)
# 输出:
# 1
# 4
# 2
# 5
# 3
# 6

ndenumerate 和 ndindex

ndenumerate

ndenumerate 提供了数组元素的索引和值。

arr = np.array([[1, 2, 3], [4, 5, 6]])
for index, value in np.ndenumerate(arr):
    print(index, value)
# 输出:
# (0, 0) 1
# (0, 1) 2
# (0, 2) 3
# (1, 0) 4
# (1, 1) 5
# (1, 2) 6

ndindex

ndindex 生成特定形状数组的多维索引。

for index in np.ndindex(arr.shape):
    print(index, arr[index])
# 输出:
# (0, 0) 1
# (0, 1) 2
# (0, 2) 3
# (1, 0) 4
# (1, 1) 5
# (1, 2) 6

flat 属性

flat 属性返回一个数组的扁平迭代器。

arr = np.array([[1, 2], [3, 4]])
for element in arr.flat:
    print(element)
# 输出:
# 1
# 2
# 3
# 4

apply_along_axis

apply_along_axis 函数允许将特定函数应用于数组的特定轴。

def my_func(x):
    return x.sum()

arr = np.array([[1, 2, 3], [4, 5, 6]])
result = np.apply_along_axis(my_func, axis=0, arr=arr)
print(result)  # 输出: [5 7 9]

result = np.apply_along_axis(my_func, axis=1, arr=arr)
print(result)  # 输出: [ 6 15]

广播迭代

使用 np.nditer 进行广播迭代：

arr_a = np.array([1, 2, 3])
arr_b = np.array([[4, 5, 6], [7, 8, 9]])
for x, y in np.nditer([arr_a, arr_b]):
    print(f"{x}, {y}")
# 输出:
# 1, 4
# 2, 5
# 3, 6
# 1, 7
# 2, 8
# 3, 9

通过这些方法，您可以灵活地遍历和操作 NumPy 数组的元素，以满足各种数据处理需求。

评论