Return value from list according to index number
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
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
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