首页 > 解决方案 > 在 Python/Pandas 中,将自定义函数应用于输入包含字符串的数据框列的最有效方法是什么?

问题描述

我有一个非常大的数据框,其中一列包含数字,另一列包含文本。我想以最有效的方式基于数字列和文本列以及复杂的自定义函数创建第三列。

根据这个来源,最有效的方法是使用 NumPy 向量化。

(下面是简化的示例代码,以阐明我尝试了什么以及卡在哪里。实际的自定义函数非常复杂,但确实将数字列和文本列作为输入。通过下面的简化代码,我想了解如何应用函数将字符串作为整个列的输入)

这完美无瑕,到目前为止一切顺利:

def fun_test1(no1, no2):
    res = no1 + no2
    return res

Test1 = pd.DataFrame({'no1':[1, 2, 3],
                     'no2':[1, 2, 3]})

Test1['result'] = fun_test1(Test1['no1'].values, Test1['no2'].values)

    no1 no2 result
0   1   1   2
1   2   2   4
2   3   3   6

然而,这不起作用,这就是我卡住的地方:

def fun_test2(no1, text):
    if text == 'one':
        no2 = 1
    elif text == 'two':
        no2 = 2
    elif text == 'three':
        no2 = 3
    res = no1 + no2
    return res

Test2 = pd.DataFrame({'no1':[1, 2, 3],
                      'text':['one', 'two', 'three']})

Test2['result'] = fun_test2(Test2['no1'].values, Test2['text'].values)

ValueError                                Traceback (most recent call last)
<ipython-input-30-a8f100d7d4bd> in <module>()
----> 1 Test2['result'] = fun_test2(Test2['no1'].values, Test2['text'].values)

<ipython-input-27-8347aa91d765> in fun_test2(no1, text)
      1 def fun_test2(no1, text):
----> 2     if text == 'one':
      3         no2 = 1
      4     elif text == 'two':
      5         no2 = 2

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

我尝试了更多变体,但最终我无法让 NumPy 矢量化处理字符串输入。

我究竟做错了什么?

如果 NumPy 向量化不适用于字符串,那么下一个最有效的方法是什么?

标签: pythonpandasperformancenumpyvectorization

解决方案


def fun_test2(no1, text, idx):
    if text[idx] == 'one':
        no2 = 1
    elif text[idx] == 'two':
        no2 = 2
    elif text[idx] == 'three':
        no2 = 3
    res = no1[idx] + no2
    return res

Test2 = pd.DataFrame({'no1':[1, 2, 3],
                      'text':['one', 'two', 'three']})

Test2['result'] = [fun_test2(Test2['no1'].values, Test2['text'].values, i) for i in range(Test2.shape[0])]

输出:

>>> Test2
   no1   text  result
0    1    one       2
1    2    two       4
2    3  three       6

或返回具有相同输出的传统方式:

def fun_test2(no1, text):
    if text == 'one':
        no2 = 1
    elif text == 'two':
        no2 = 2
    elif text == 'three':
        no2 = 3
    res = no1 + no2
    return res

Test2 = pd.DataFrame({'no1':[1, 2, 3],
                      'text':['one', 'two', 'three']})

Test2['result'] = [fun_test2(Test2['no1'].values[i], Test2['text'].values[i]) for i in range(Test2.shape[0])]

推荐阅读