Создайте новый столбец pd.DataFrame, если значение существующего столбца содержит указанную c подстроку

Question

У меня есть DataFrame, как показано ниже:

     entity   
0     cat  
1     dog
2     tree
3     dog
4     flower

Я хочу создать новый столбец с именем type, который содержит значения, основанные на условии, что значение в entity содержит подстроку

substring1 = 'cat|og'
substring2 = 'ree|ower'

результирующий DataFrame должен выглядеть следующим образом:

    entity      type   
0     cat      animal
1     dog      animal
2     tree     plant
3     dog      animal
4     flower   plant

Подход, который я использую, это взять pd.Series.str.contains и создать два логических столбца is_animal и is_plant.

Как я могу изменить приведенный ниже код, чтобы устранить необходимость в двух логических столбцах и просто иметь столбец type?

df.loc[:, 'is_animal'] = df.entity.astype(str).str.contains(substring1, case=False)
df.loc[:, 'is_plant'] = df.entity.astype(str).str.contains(substring2, case=False)

Я попытался перебрать DataFrame используя iteritems() для создания столбца type, но, похоже, он не работает.

Answer 1

Вы можете создать словарь для type s по подстановкам, а затем установить новые значения путем сравнения substring s только для соответствующих строк:

d = {'animal':substring1, 'plant':substring2}

s = df.entity.astype(str)
for k, v in d.items():
    df.loc[s.str.contains(v, case=False), 'type'] = k
print (df)
   entity    type
0     cat  animal
1     dog  animal
2    tree   plant
3     dog  animal
4  flower   plant