Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Group pandas DataFrame by columns and produce lists out of groups of columns

Assuming we have the following input DataFrame:

In [80]: %paste
data = {
 'Item_2_id': {0: 24, 1: 41, 2: 34},
 'Item_2_quantity': {0: 4, 1: 1, 2: 4},
 'Item_3_id': {0: 16, 1: 33, 2: 8},
 'Item_3_quantity': {0: 1, 1: 1, 2: 2},
 'customer_name': {0: 'John', 1: 'Paul', 2: 'Andrew'},
 'item_1_id': {0: 4, 1: 8, 2: 1},
 'item_1_quantity': {0: 1, 1: 3, 2: 1},
 'order_id': {0: 1, 1: 2, 2: 3}
}
cols = 'order_id customer_name  item_1_id  item_1_quantity  Item_2_id  Item_2_quantity  Item_3_id  Item_3_quantity'.split()

df = pd.DataFrame(data)[cols]

df
## -- End pasted text --
Out[80]:
   order_id customer_name  item_1_id  item_1_quantity  Item_2_id  Item_2_quantity  Item_3_id  Item_3_quantity
0         1          John          4                1         24                4         16                1
1         2          Paul          8                3         41                1         33                1
2         3        Andrew          1                1         34                4          8                2

How can we group all id and quantity columns, so that we will get the following desired DataFrame:

In [85]: result
Out[85]:
   order_id customer_name           id   quantity
0         1          John  [4, 24, 16]  [1, 4, 1]
1         2          Paul  [8, 41, 33]  [3, 1, 1]
2         3        Andrew   [1, 34, 8]  [1, 4, 2]

My attempts:

In [191]: id_vars = ['order_id','customer_name']

In [192]: df.set_index(id_vars) \
            .groupby(lambda x: x.split('_')[-1], axis=1) \
            .agg(lambda x: x.tolist())
Out[192]:
                                                 id                                       quantity
order_id customer_name
1        John           (i, t, e, m, _, 1, _, i, d)  (i, t, e, m, _, 1, _, q, u, a, n, t, i, t, y)
2        Paul           (I, t, e, m, _, 2, _, i, d)  (I, t, e, m, _, 2, _, q, u, a, n, t, i, t, y)
3        Andrew         (I, t, e, m, _, 3, _, i, d)  (I, t, e, m, _, 3, _, q, u, a, n, t, i, t, y)

if i just print it - it works properly:

In [193]: df.set_index(id_vars) \
            .groupby(lambda x: x.split('_')[-1], axis=1) \
            .agg(lambda x: print(x.tolist()))
[4, 24, 16]
[8, 41, 33]
[1, 34, 8]
[1, 4, 1]
[3, 1, 1]
[1, 4, 2]
Out[193]:
                          id quantity
order_id customer_name
1        John           None     None
2        Paul           None     None
3        Andrew         None     None

PS actually I came to that problem when I was trying to answer another question and i found another solution, but I feel like there must be much more elegant solution, which is using something like:

df.groupby(..., axis=1).agg(...)

or

df.groupby(..., axis=1).apply(...)
like image 473
MaxU - stop WAR against UA Avatar asked Mar 09 '23 11:03

MaxU - stop WAR against UA


2 Answers

using filter

df[['order_id', 'customer_name']].assign(
    id=df.filter(regex='[Ii]tem_\d+_id').values.tolist(),
    quantity=df.filter(regex='[Ii]tem_\d+_quantity').values.tolist()
)

   order_id customer_name           id   quantity
0         1          John  [4, 24, 16]  [1, 4, 1]
1         2          Paul  [8, 41, 33]  [3, 1, 1]
2         3        Andrew   [1, 34, 8]  [1, 4, 2]

enter image description here

like image 185
piRSquared Avatar answered Apr 06 '23 00:04

piRSquared


Here is a very nice solution from @DSM:

In [123]: df.set_index(['order_id','customer_name']) \
     ...:   .groupby(lambda x: x.split('_')[-1], axis=1) \
     ...:   .agg(lambda x: x.values.tolist())
     ...:
Out[123]:
                                 id   quantity
order_id customer_name
1        John           [4, 24, 16]  [1, 4, 1]
2        Paul           [8, 41, 33]  [3, 1, 1]
3        Andrew          [1, 34, 8]  [1, 4, 2]
like image 22
MaxU - stop WAR against UA Avatar answered Apr 05 '23 23:04

MaxU - stop WAR against UA