Pandas小册子：根据条件创建新的列

在进行数据分析时，经常会遇到根据已有的数据列，按照一定条件创建新的数据列，然后进行进一步分析。

今天，我们来看一个根据已有数据按照一定条件创建新的数据列的方法。

数据如下：

1. import pandas as pd

3. df = pd.DataFrame({'team_A': ['Spain', 'Germany', 'Brazil', 'France'],

4.                  'team_B': ['USA', 'Argentina', 'Mexico', 'Belgium'],

5.                  'score_A': [5, 3, 2, 0],

6.                  'score_B': [4, 0, 3, 0]},

7.                 columns = ['team_A', 'team_B', 'score_A', 'score_B'])

8. df

10. Out[2]:

11.    team_A     team_B  score_A  score_B

12. 0    Spain        USA        5        4

13. 1  Germany  Argentina        3        0

14. 2   Brazil     Mexico        2        3

15. 3   France    Belgium        0        0

问题： 从上面数据中创建新的一个数据列，用来存储获胜队伍的名称。

即，根据 “score_A” 与 “score_B” 比较的结果，来获取相应的结果。

例如，第一行， “Spain”:”USA” 为 5:4 ， “Spain” 获胜， 新创建的列中存储的结果为 “Spain”。

下面介绍两种方法来实现上述要求。

第一种方法

第一种方法是 利用 Pandas 中 DataFrame 的条件选择功能来实现，过程如下：

1. # 创建新的列 "win_team"，赋值为空白

2. df['win_team'] = ''

4. # 创建判断条件 mask

5. mask = df['score_A'] - df['score_B']

6. df.loc[mask > 0, 'win_team'] = df.loc[mask > 0, 'team_A']

7. df.loc[mask < 0, 'win_team'] = df.loc[mask < 0, 'team_B']

8. df.loc[mask == 0, 'win_team'] = 'Draw'

10. df

12. Out[3]:

13.    team_A     team_B  score_A  score_B win_team

14. 0    Spain        USA        5        4    Spain

15. 1  Germany  Argentina        3        0  Germany

16. 2   Brazil     Mexico        2        3   Mexico

17. 3   France    Belgium        0        0     Draw

第二种方法

第二种方法是结合 DataFrame.iterrows() 以及 Python 的 list 的功能来实现，过程如下：

1. # The second method to get the winners

3. def find_win_team(df):

4.    winners = []

5.    for i, row in df.iterrows():

6.        if row['score_A'] > row['score_B']:

7.            winners.append(row['team_A'])

8.        elif row['score_A'] < row['score_B']:

9.            winners.append(row['team_B'])

10.        else:

11.            winners.append('Draw')

12.    return winners

14. df['winner'] = find_win_team(df)

15. df

17. Out[4]:

18.    team_A     team_B  score_A  score_B win_team   winner

19. 0    Spain        USA        5        4    Spain    Spain

20. 1  Germany  Argentina        3        0  Germany  Germany

21. 2   Brazil     Mexico        2        3   Mexico   Mexico

22. 3   France    Belgium        0        0     Draw     Draw

关于 DataFrame.iterrows(), 我们先来看看其运行结果。

1. for row_index, row in df.iterrows():

2.    print('%s\n%s' % (row_index, row))

4. 0

5. team_A      Spain

6. team_B        USA

7. score_A         5

8. score_B         4

9. win_team    Spain

10. winner      Spain

11. Name: 0, dtype: object

12. 1

13. team_A        Germany

14. team_B      Argentina

15. score_A             3

16. score_B             0

17. win_team      Germany

18. winner        Germany

19. Name: 1, dtype: object

20. 2

21. team_A      Brazil

22. team_B      Mexico

23. score_A          2

24. score_B          3

25. win_team    Mexico

26. winner      Mexico

27. Name: 2, dtype: object

28. 3

29. team_A       France

30. team_B      Belgium

31. score_A           0

32. score_B           0

33. win_team       Draw

34. winner         Draw

35. Name: 3, dtype: object

fDataFrame.iterrows() 的作用是将 dataframe的每行转换成为一个 Series，可以理解为 针对于每一行，做了行列转置。

图示如下：