Return column of pandas df when another column's list contains specific value

user_15309695 0
Viewed 275

I have a df constructed as such:

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)

The resulting df looks like

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

I would like to write a function that returns a dataframe or series of ids if a specific value in the corresponding list was called. For example:

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

So

find_ids('two')

should return

id
001
002

and

find_ids('twelve')

should return

id
003
pandas list python
3 Answers

Try .str.join(sep = " ").str.contains(value) which first converts the list into string and then checks if generated string contains the value.

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

Output:

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

To make it more efficient try saving list as string using .str.join(sep = " ") in a new column and then you can search using .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"]

Output:

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

You can try:

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


print(find_ids(df, "two"))

Prints:

0    001
1    002
Name: id, dtype: object
user_100359850

You can use:

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

Now finally call the function:

print(find_ids('two'))

Output:

0    001
1    002
user_142898920
Related