The pandas DataFrame group by column and lists lists of column groups

Question

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 the columns idand quantityso that we 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 correctly:

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 this problem when I tried to answer another question and I found a different solution, but I feel that there should be a much more elegant solution that uses something like:

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

or

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

Answer 1

via 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]

---

Answer 2

@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]

Answer 3

, @MaxU:

id_cols = [col for col in list(df) if 'id' in col and col not in 'order_id']
quant_cols = [col for col in list(df) if 'quantity' in col]

df2 = df[['order_id', 'customer_name']].copy()
df2['quantity'] = list(df[quant_cols].values)
df2['id'] = list(df[id_cols].values)

df2:
   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]