Pandas: groupby to list

user_8236090 0
Viewed 21784

I have data like below:

id  value   time

1   5   2000
1   6   2000
1   7   2000
1   5   2001
2   3   2000
2   3   2001
2   4   2005
2   5   2005
3   3   2000
3   6   2005

My final goal is to have data in a list like below:

[[5,6,7],[5]] (this is for id 1 grouped by the id and year)
[[3],[3],[4,5]] (this is for id 2 grouped by the id and year)
[[3],[6]] (same logic as above)

I have grouped the data using df.groupby(['id', 'year']). But after that, I am not able to access the groups and get the data in the above format.

pandas python
3 Answers

You can use apply(list):

>>> df.groupby(['id', 'time'])['value'].apply(list)

id  time
1   2000    [5, 6, 7]
    2001          [5]
2   2000          [3]
    2001          [3]
    2005       [4, 5]
3   2000          [3]
    2005          [6]
Name: value, dtype: object

If you really want it in the exact format as you displayed, you can then groupby id and apply list again, but this is not efficient, and that format is arguably harder to work with...

>>> df.groupby(['id','time'])['value'].apply(list).groupby('id').apply(list).tolist()
[[[5, 6, 7], [5]], [[3], [3], [4, 5]], [[3], [6]]]
user_66711760

If you want to calculate the lists for multiple columns, you can do the following:

import pandas as pd

df = pd.DataFrame(
    {'A': [1,1,2,2,2,2,3],
     'B':['a','b','c','d','e','f','g'],
     'C':['x','y','z','x','y','z','x']})

df.groupby('A').agg({'B': list,'C': list})

Which will calculate lists of B and C:

              B             C
A                            
1        [a, b]        [x, y]
2  [c, d, e, f]  [z, x, y, z]
3           [g]           [x]

To get lists for all columns:

df.groupby('A').agg(list)

To have the lists be sorted:

df.groupby('A').agg(sorted)
user_10603490

You could do the following:

import pandas as pd

data = [[1, 5, 2000],
        [1, 6, 2000],
        [1, 7, 2000],
        [1, 5, 2001],
        [2, 3, 2000],
        [2, 3, 2001],
        [2, 4, 2005],
        [2, 5, 2005],
        [3, 3, 2000],
        [3, 6, 2005]]

df = pd.DataFrame(data=data, columns=['id', 'value', 'year'])

result = []
for name, group in df.groupby(['id']):
    result.append([g['value'].values.tolist() for _, g in group.groupby(['year'])])

for e in result:
    print(e)

Output

[[5, 6, 7], [5]]
[[3], [3], [4, 5]]
[[3], [6]]
user_40015920
Related