2018年4月20日 星期五

Python - 使用apply與自定義function建立dataframe的新欄位,常用於分組 - Applying function with multiple arguments to create a new pandas column

Information:

System version : Windows 10 64-bit
Python version : Python 3.6.0 :: Anaconda 4.3.1 (64-bit)

Code:

import pandas as pd
import numpy as np

df1 = pd.DataFrame()
number = [1,2,3,4,5]
sex = ['male','male','female','female','female']
df1['number'] = number
df1['sex'] = sex
df1['income'] = [500,1500,1300,2500,500]
print(df1)
print('-'*45)

def income_group(df):
    if df['income']>2000:
        return 'high'
    elif df['income']>1000:
        return 'medium'
    else:
        return 'low'

df1['group'] = df1.apply(income_group, axis=1)
print(df1)
print('-'*45)

Result:

   number     sex  income
0       1    male     500
1       2    male    1500
2       3  female    1300
3       4  female    2500
4       5  female     500
---------------------------------------------
   number     sex  income   group
0       1    male     500     low
1       2    male    1500  medium
2       3  female    1300  medium
3       4  female    2500    high
4       5  female     500     low
---------------------------------------------

沒有留言:

張貼留言