如何在 NumPy 数组上映射函数？

在进行科学计算和数据分析时，NumPy 数组是常用的数据结构。在数组上进行函数映射是常见的操作之一。本文将介绍如何在 NumPy 数组上映射函数。

使用 np.vectorize 函数

在 NumPy 中，可以使用 np.vectorize 函数对一般函数进行向量化操作，进而在数组上进行函数映射。

import numpy as np

def my_func(x):
    return x ** 2 + 1

x = np.array([1, 2, 3])
y = np.vectorize(my_func)(x)
print(y)
# 输出: [ 2  5 10]

需要注意的是，向量化操作并不会比普通循环快很多，因为其仍然需要进行逐个元素计算。

使用 np.apply_along_axis 函数

如果需要在数组的某个轴上进行函数映射，可以使用 np.apply_along_axis 函数。

import numpy as np

def my_func(x):
    return x ** 2 + 1

x = np.array([[1, 2], [3, 4]])
y = np.apply_along_axis(my_func, 0, x)
print(y)
# 输出: [[ 2  5]
#       [10 17]]

在这个例子中，np.apply_along_axis 函数将 my_func 函数应用于 x 数组的轴 0（即行）上，返回新的数组 y。

使用数组广播

NumPy 数组的广播功能可以使得形状不同的数组进行逐元素计算。在一些情况下，可以利用数组广播来进行函数映射。

例如，可以将函数向量化并利用广播机制进行点乘操作：

import numpy as np

def my_func(x):
    return x ** 2 + 1

x = np.array([1, 2, 3])
y = my_func(x)
z = x * y
print(z)
# 输出: [ 2 10 28]

需要注意的是，这种方法只适用于能够进行逐元素计算的情况，如加减乘除等。

总结

本文介绍了三种在 NumPy 数组上进行函数映射的方法，分别是：

• 使用 np.vectorize 函数
• 使用 np.apply_along_axis 函数
• 利用数组广播机制

在实际编程中，可以根据具体情况选择合适的方法。