Return value from list according to index number

user_6366770 0
Viewed 213

I have been struggling to find the pandas solution to this without looping:

input:

df = pd.DataFrame({'A' : [[6,1,1,1], [1,5,1,1], [1,1,11,1], [1,1,1,20]]})

                A
0   [6, 1, 1, 1]
1   [1, 5, 1, 1]
2   [1, 1, 11, 1]
3   [1, 1, 1, 20]

output:

                A   B
0   [6, 1, 1, 1]    6
1   [1, 5, 1, 1]    5
2   [1, 1, 11, 1]   11
3   [1, 1, 1, 20]   20

I have tried so many different things over the past hour or so, and I know the solution will be an embarrassingly simple one-liner. Thanks for your help -- not my python day today!

pandas python
4 Answers

You can do a simple list comprehension:

df['B'] = [s[i] for i, s in zip(df.index, df['A'])]

Or if you want only diagonal values:

df['B'] = np.diagonal([*df['A']])

               A   B
0   [6, 1, 1, 1]   6
1   [1, 5, 1, 1]   5
2  [1, 1, 11, 1]  11
3  [1, 1, 1, 20]  20
user_128331660

Solution using numpy:

import pandas as pd 
import numpy as np

df = pd.DataFrame({'A' : [[6,1,1,1], [1,5,1,1], [1,1,11,1], [1,1,1,20]]})
df['B'] = pd.Series(np.diag(np.vstack(df['A'].to_numpy())), index=df.index)

df
user_141686230

I have also been able to solve my own question, but terribly. I will post the solution anyway:

df = pd.DataFrame({'A' : [[6,1,1,1], [1,5,1,1], [1,1,11,1], [1,1,1,20]]})
s = df.explode('A')
i = s.groupby(level=0).cumcount()
df['B'] = s.loc[s.index == i, 'A']

               A   B
0   [6, 1, 1, 1]   6
1   [1, 5, 1, 1]   5
2  [1, 1, 11, 1]  11
3  [1, 1, 1, 20]  20
user_63667700

You can use apply() and the dataframe index (row.name) to extract the diagonal elements.

import pandas as pd
df = pd.DataFrame({'A' : [[6,1,1,1], [1,5,1,1], [1,1,11,1], [1,1,1,20]]})

df['B']=df.apply(lambda row: row['A'][row.name],axis=1)

print(df)

Output :

    A              B
0   [6, 1, 1, 1]   6
1   [1, 5, 1, 1]   5
2  [1, 1, 11, 1]  11
3  [1, 1, 1, 20]  20
user_135029120
Related