Возвращает столбец pandas df, когда список другого столбца содержит определенное значение

#python #pandas #list

Вопрос:

У меня есть df, построенный как таковой:

 import pandas as pd

dic = {'001': [['one','two','three']],
       '002': [['two', 'five', 'eight']],
       '003': [['three','six','ten','twelve']]}
df = pd.DataFrame.from_dict(dic,orient='index')
df.reset_index(inplace=True)
df = df.rename(columns = {'index':'id',0:'values'})
print(df)
 

Полученный df выглядит следующим образом

     id                     values
0  001          [one, two, three]
1  002         [two, five, eight]
2  003  [three, six, ten, twelve]
 

Я хотел бы написать функцию, которая возвращает фрейм данных или серию идентификаторов, если было вызвано определенное значение в соответствующем списке. Например:

 def find_ids(value):
    ids = psuedocode: if list contains value, then return id
    return ids
 

Так

 find_ids('two')
 

должен вернуться

 id
001
002
 

и

 find_ids('twelve')
 

должен вернуться

 id
003
 

Ответ №1:

Попробуйте .str.join(sep = " ").str.contains(value) сначала преобразовать list в string , а затем проверить, содержит ли сгенерированная строка value .

 def find_ids(df, value):
   return df.loc[df['values'].str.join(sep = " ").str.contains(value), "id"]
 

Выход:

 >>> print(find_ids(df, "two"))
0    001
1    002
Name: id, dtype: object
 

Чтобы сделать его более эффективным, попробуйте сохранить список в виде строки, используя .str.join(sep = " ") в новом столбце, а затем вы можете выполнить поиск с помощью .str.contains(value)

 df['values_str'] = df['values'].str.join(sep = " ")
def find_ids(df, value):
    return df.loc[df.values_str.str.contains(value), "id"]
 

Выход:

 >>> print(find_ids(df, "two"))
0    001
1    002
Name: id, dtype: object
 

Ответ №2:

Вы можете попробовать:

 def find_ids(df, value):
    return df.loc[df["values"].apply(lambda x: value in x), "id"]


print(find_ids(df, "two"))
 

С принтами:

 0    001
1    002
Name: id, dtype: object
 

Ответ №3:

Вы можете использовать:

 def find_ids(value):
    newdf=df.explode('values')
    return newdf.loc[newdf['values']==value,'id']
 

Теперь, наконец, вызовите функцию:

 print(find_ids('two'))
 

Выход:

 0    001
1    002