extract non-nested list from pandas dataframe

Viewed 121
df = pd.DataFrame({'x':['a','b','c','d'],'y':[[1],[3,4],[2],[5]]})

I want to extract the list [3,4] so I did

df.loc[df.x=='b','y'].tolist()

but it returns a nested list.

[[3,4]]

How can I get a non-nested list from the dataframe?

4 Answers

You can explode before

df1=df.explode('y')
df1.loc[df1.x=='b','y'].tolist()
... 
[3, 4]

Also simple fix

df.loc[df.x=='b','y']
[3, 4]

Your loc expression returns a sequence rather than a scalar, you just need to select its first and only element:

df.loc[df.x=='b','y'].iloc[0]

If you're trying to get a single cell value, try df.at, it guarantees a single value:

df.at[pd.Index(df['x']).get_loc('b'), 'y']
# [3, 4]

This solution is guaranteed to fail if the key does not exist:

df.at[pd.Index(df['x']).get_loc('xyz'), 'y']
# KeyError: 'xyz'

The pd.Index(df['x']).get_loc('b') expression could also be replaced with df['x'].eq(b).argmax(), but be warned it will return an incorrect result if "b" does not exist.

Try loc and squeeze

In [621]: df.loc[df.x=='b','y'].squeeze()
Out[621]: [3, 4]
Related