首页 > 解决方案 > Numpy将函数应用于数组中的每个项目

问题描述

所以假设我有一个二维数组。如何将函数应用于数组中的每个项目并用返回替换该项目?此外,函数的返回将是一个元组,因此数组将变为 3d。

这是记住的代码。

def filter_func(item):
    if 0 <= item < 1:
        return (1, 0, 1)
    elif 1 <= item < 2:
        return (2, 1, 1)
    elif 2 <= item < 3:
        return (5, 1, 4)
    else:
        return (4, 4, 4)

myarray = np.array([[2.5, 1.3], [0.4, -1.0]])

# Apply the function to an array

print(myarray)

# Should be array([[[5, 1, 4],
#                   [2, 1, 1]],
#                  [[1, 0, 1],
#                   [4, 4, 4]]])

任何想法我该怎么做?一种方法是这样做np.array(list(map(filter_func, myarray.reshape((12,))))).reshape((2, 2, 3)),但这很慢,尤其是当我需要在形状数组(1024、1024)上进行时。

我也看到人们使用 np.vectorize,但它以某种方式结束,(array([[5, 2], [1, 4]]), array([[1, 1], [0, 4]]), array([[4, 1], [1, 4]])).然后它的形状为 (3, 2, 2)。

标签: pythonarraysnumpy

解决方案


无需更改函数中的任何内容。

只需将函数的矢量化版本应用于数组并堆叠结果:

np.stack(np.vectorize(filter_func)(myarray), axis=2)

结果是:

array([[[5, 1, 4],
        [2, 1, 1]],

       [[1, 0, 1],
        [4, 4, 4]]])

推荐阅读