沿一个轴应用 np.vectorize

Question

假设我有两个数组 arr1 和 arr2:

arr1 = [0, 1, 2]

arr2 = [
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8],
]

假设我有一个函数可以对这个数组的元素做一些事情：

def func(arr):
    new_arr = arr.copy()
    new_arr[0] = new_arr[0] * 2
    new_arr[1] = new_arr[1] * 10
    new_arr[2] = new_arr[2] * 100
    return new_arr

现在我想对其进行矢量化，以便它适用于 arr1 和 arr2:

func(arr1)
# returns [0, 10, 200]
func(arr2)
# returns
# [0, 10, 200],
# [6, 40, 500],
# [12, 70, 800],

np.vectorize 不起作用，因为它分解了数组参数中的每个元素。我希望它仅沿第一个轴应用该功能。

np.apply_along_axis 几乎可以工作，除了它不会将一维数组参数视为单个参数。

最好的方法是什么？

Answer 1

你可以直接将数组相乘。感谢 numpy broadcasting:

factor = np.array([2, 10, 100])

arr1 * factor
array([  0,  10, 200])

arr2 * factor
array([[  0,  10, 200],
       [  6,  40, 500],
       [ 12,  70, 800]])

Answer 2

只需这样做：

import numpy as np

a = np.array([0, 1, 2])

b = np.array([
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8],
])

c = np.array([2, 10, 100])

print(a*c)

print(b*c)

输出：

[  0  10 200]
[[  0  10 200]
 [  6  40 500]
 [ 12  70 800]]

Answer 3

如果您花时间阅读 np.vectorize 文档，您最终会遇到 signature 选项：

In [27]: f= np.vectorize(func, signature='(n)->(n)')
In [28]: f(arr1)
Out[28]: array([  0,  10, 200])
In [29]: f(arr2)
Out[29]: 
array([[  0,  10, 200],
       [  6,  40, 500],
       [ 12,  70, 800]])

进一步阅读您会遇到关于性能的警告。

沿一个轴应用 np.vectorize

Apply np.vectorize along one axis

python

numpy

vectorization