将函数应用于python中的数据框

Question

在以下代码中，我尝试先定义一个函数并将该函数应用于数据框以重置地理区域。

import pandas as pd
testdata ={'country': ['USA','AUT','CHE','ABC'], 'geozone':[0,0,0,0]}
d =pd.DataFrame.from_dict(testdata, orient = 'columns')
    
def setgeozone(dataframe, dcountry, dgeozone):
    dataframe.loc[dataframe['dcountry'].isin(['USA','CAN']),'dgeozone'] =1
    dataframe.loc[dataframe['dcountry'].isin(['AUT','BEL']),'dgeozone'] =2
    dataframe.loc[dataframe['dcountry'].isin(['CHE','DNK']),'dgeozone'] =3
        
setgeozone(d, country, geozone)

我收到错误消息：

Traceback (most recent call last):

  File "<ipython-input-56-98dad4781f73>", line 1, in <module>
    setgeozone(d, country, geozone)

NameError: name 'country' is not defined

谁能帮我理解我做错了什么。

非常感谢。

Answer 1

您不需要将 DataFrame 本身以外的参数传递给您的函数。试试这个：

def setgeozone(df):
    df.loc[df['country'].isin(['USA','CAN']),'geozone'] = 1
    df.loc[df['country'].isin(['AUT','BEL']),'geozone'] = 2
    df.loc[df['country'].isin(['CHE','DNK']),'geozone'] = 3

setgeozone(df)

这里有另外两种（也是更好的）方法来完成你需要的：

1. 使用map:

df["geozone"] = df["country"].map({"USA": 1, "CAN": 1, "AUT": 2, "BEL": 2, "CHE": 3, "DNK": 3})

1. 使用numpy.select:

import numpy as np
df["geozone"] = np.select([df["country"].isin(["USA", "CAN"]), df["country"].isin(["AUT", "BEL"]), df["country"].isin(["CHE", "DNK"])], 
                          [1, 2, 3])