Python - stuck in a loop when trying to store an array of list

Question

I am iterating through the dataframe and try to store the "total" list for each id.

The dataframe look like this:

id    total    difference
 1     29         3
 1     21         2
 2     39         0
 2     22         9

What I have tried so far:

total_list=[]
for i, row in df.iterrows(): 
    total_list.extend(df.total.loc[df.id==row.id].tolist())

print(total_list) # this never gets print

total_list should look like [[29, 21], [39,22]]

Answer 1

This can be done without looping, as long as I understand correctly. I'm assuming that 31 in your output should be a 21.

Setup

>>> df = pd.DataFrame([[1,29,3],[1,21,2],[2,39,0],[2,22,9]], columns=['id', 'total','difference'])
>>>
>>> df
   id  total  difference
0   1     29           3
1   1     21           2
2   2     39           0
3   2     22           9

Solution

>>> df.groupby('id')['total'].apply(list).tolist()
[[29, 21], [39, 22]]