遍歷if語句的行並添加到新列[Pandas / Python]

Question

我有一個3列的數據框

           Hospital              2009-10  2010-11
Aberystwyth Mental Health Unit     19       19
Bro Ddyfi Community Hospital       16       10  
Bronglais General Hospital         160      148
Caebryn Mental Health Unit         37       39 
Carmarthen Mental Health Unit      38       31  

我正在嘗試創建一個函數來檢查單詞是否在醫院列中，如果這樣，它將單詞放入新列中，如下所示：

               Hospital              2009-10  2010-11    Hospital Type
    Aberystwyth Mental Health Unit     19       19          Mental
    Bro Ddyfi Community Hospital       16       10          Community  
    Bronglais General Hospital         160      148         General
    Caebryn Mental Health Unit         37       39          Mental
    Carmarthen Mental Health Unit      38       31          Mental

這是我嘗試過的代碼：

def find_type(x):
    if df['Hospital'].str.contains("Mental").any():
        return "Mental"
    if df['Hospital'].str.contains("Community").any():
        return "Community"
    else:
        return "Other"

df['Hospital Type'] = df.apply(find_type)

我得到的輸出是這樣的：

               Hospital              2009-10  2010-11    Hospital Type
    Aberystwyth Mental Health Unit     19       19          NaN
    Bro Ddyfi Community Hospital       16       10          NaN  
    Bronglais General Hospital         160      148         NaN
    Caebryn Mental Health Unit         37       39          NaN
    Carmarthen Mental Health Unit      38       31          NaN

我怎樣才能得到預期的輸出呢？

謝謝！

Answer 1

使用以|分隔的關鍵字extract 與fillna ：

pat = r"(Mental|Community)"
df['Hospital Type'] = df['Hospital'].str.extract(pat, expand=False).fillna('Other')
print (df)
                         Hospital  2009-10  2010-11 Hospital Type
0  Aberystwyth Mental Health Unit       19       19        Mental
1    Bro Ddyfi Community Hospital       16       10     Community
2      Bronglais General Hospital      160      148         Other
3      Caebryn Mental Health Unit       37       39        Mental
4   Carmarthen Mental Health Unit       38       31        Mental