How to create a list based on 2 conditions from a dataframe?-CodePudding

I have this dataframe :

 A       B      C
 1       1     100
 1       1     101
 1       5     102
 2       6     103
 2       6     104
 3      19     105
 3      20     106
...    ...      ...
 n

From each index of 'A' AND each index of the column 'B' I would like to obtain a list that respects the 2 conditions based on 'C' values for the whole dataframe.

For example, I would like this output:

l1 = [100, 101]
l2 = [103,104]
l3 = [105, 106]
etc.

Can you help please ? Thanks !

CodePudding user response：

You could use the pd.Series slices to get just the first two elements in a list -

df.groupby('A')['C'].apply(lambda x: list(x[:2]))

Output

A
1    [100, 101]
2    [103, 104]
3    [105, 106]

CodePudding user response：

If need first 2 rows per groups by A column use:

s = df.groupby('A').head(2).groupby('A')['C'].agg(list)

#alternative solution
#s = df.groupby('A')['C'].agg(lambda x: list(x.iloc[:2]))
print (s)
A
1    [100, 101]
2    [103, 104]
3    [105, 106]
Name: C, dtype: object


d = df.groupby('A').head(2).groupby('A')['C'].agg(list).to_dict()
print (d)
{1: [100, 101], 2: [103, 104], 3: [105, 106]}

print (d[1])
[100, 101]