Custom Melt / Collapse Pandas

user_7987118 0
Viewed 81

I have another problem with pandas. I can do the below task utilizing loops but it would be very inefficient due to the size of the input. Please let me know if there is a pandas solution.

I would like create a new DF_C based on DF A. I need to create multiple rows based on the columns COL_A and COL_B (the values will be separated by commas). State will always have one element in it.

The sequence of rows does not matter.

I have a DF A:

State       COL_A        B_COL
01           01          03, 01
02           01, 03      01, 04
02           07          03
04           01          05

I would like a resulting df_c:

 State       COL_A        B_COL
 01            01          03
 01            01          01
 02            01          01
 02            01          04
 02            03          01 
 02            03          04
 02            07          03
 04            01          05
pandas numpy python
2 Answers

you can do by first use str.split on both COL_A and B_COL, then chain with one explode on each column like:

df_ = (df.assign(COL_A=lambda x: x['COL_A'].str.split(', '), 
                 B_COL=lambda x: x['B_COL'].str.split(', '))
         .explode('COL_A')
         .explode('B_COL')
       )
print (df_)
   State COL_A B_COL
0      1    01    03
0      1    01    01
1      2    01    01
1      2    01    04
1      2    03    01
1      2    03    04
2      2    07    03
3      4    01    05

EDIT: if you are after efficiency, maybe consider doing

df_ = pd.DataFrame(
            [(s, a, b) 
             for s, cola, colb in zip(df['State'], df['COL_A'], df['B_COL'])
             for a in cola.split(', ') for b in colb.split(', ')], 
            columns=df.columns)
user_92747320

An alternative to Ben.T's second solution, using itertools : 

from itertools import product,chain

flatten = chain.from_iterable

result = flatten(product([state],col_a.split(","),b_col.split(","))
          for state, col_a,b_col in df.to_numpy())

pd.DataFrame(result, columns = df.columns)

 State  COL_A   B_COL
0   1    01     03
1   1    01     01
2   2    01     01
3   2    01     04
4   2    03     01
5   2    03     04
6   2    07     03
7   4    01     05
user_71757130
Related