如何在numpy字符串数组上实现自定义函数的向量化操作

有两种成熟的实现方式，分别是用np.frompyfunc和np.vectorize，具体用法如下：

方案1：使用np.frompyfunc（性能更优，推荐） #

np.frompyfunc是numpy底层提供的Python函数转向量化函数工具，运行效率更高：

import numpy as np

# 自定义函数，这里是取列表第一个元素
def custom(text):
    return text[0]

# 转换为向量化函数：参数依次为 待转换函数、输入参数数量、返回值数量
vec_custom = np.frompyfunc(custom, 1, 1)

test_array = np.char.array(["sample1-sample","sample2-sample"])
split_arr = test_array.split('-')
result = vec_custom(split_arr)

# 如需转换为chararray类型，可执行：result = np.char.array(result)
print(result)

输出结果：

array(['sample1', 'sample2'], dtype=object)

方案2：使用np.vectorize（语法更友好） #

如果需要指定输出类型，不用后续手动转换，可以用np.vectorize包装：

import numpy as np

def custom(text):
    return text[0]

# otypes指定输出为长度7的字符串类型，和原示例输出类型一致
vec_custom = np.vectorize(custom, otypes=['U7'])

test_array = np.char.array(["sample1-sample","sample2-sample"])
split_arr = test_array.split('-')
result = vec_custom(split_arr)

print(result)

输出结果：

array(['sample1', 'sample2'], dtype='<U7')

注意事项 #

• 两种方法都支持任意自定义Python函数的向量化转换，无需依赖numpy内置的字符串处理函数
• 性能排序：np.frompyfunc > np.vectorize > 普通Python循环，数据量较大时优先选择np.frompyfunc
• 若自定义函数逻辑可通过numpy广播、内置算子重构，性能会优于上述Python函数包装方案，逻辑过于复杂无法用内置操作实现时优先选择上述方案

内容的提问来源于stack exchange，提问作者data_person