Pandas - Align column subset rows with matching values in master column

user_16620792 0
Viewed 141

I am starting out with the following dataframe (df):

  variable  key  i  j  k
0        a    1  0  0  0
1        a    2  0  0  0
2        a    3  0  0  0
3        b    2  0  0  0
4        c    1  0  0  0
5        c    3  0  0  0

Using the following code, I transform it into the 2nd displayed dataframe (out).

out = pd.DataFrame()
out['keys_master'] = list(range(min(df['key']), max(df['key']) + 1))
for value in ['a', 'b', 'c']:
    temp = df.loc[df[variable_string] == value]
    cols = [f'{i}{value}' for i in df.columns.values]
    temp.columns = cols
    out = pd.concat([out, temp], axis=1)

    # Drop columns that contain variable names, then drop empty rows
    out.drop(columns=cols[0], inplace=True)
    out.dropna(axis=0, how='all', inplace=True)

out:

   keys_master  keya   ia   ja   ka  keyb   ib   jb   kb  keyc   ic   jc   kc
0          1.0   1.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN   NaN  NaN  NaN  NaN
1          2.0   2.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN   NaN  NaN  NaN  NaN
2          3.0   3.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN   NaN  NaN  NaN  NaN
3          NaN   NaN  NaN  NaN  NaN   2.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN
4          NaN   NaN  NaN  NaN  NaN   NaN  NaN  NaN  NaN   1.0  0.0  0.0  0.0
5          NaN   NaN  NaN  NaN  NaN   NaN  NaN  NaN  NaN   3.0  0.0  0.0  0.0

How can I transform out into a dataframe like this?

   keys_master  keya  ia  ja  ka  keyb  ib  jb  kb  keyc  ic  jc  kc
0          1.0   1.0 0.0 0.0 0.0   NaN NaN NaN NaN   1.0 0.0 0.0 0.0
1          2.0   2.0 0.0 0.0 0.0   2.0 0.0 0.0 0.0   NaN NaN NaN NaN
2          3.0   3.0 0.0 0.0 0.0   NaN NaN NaN NaN   3.0 0.0 0.0 0.0
pandas dataframe python
1 Answers

Similar to the original approach of grouping together and concating. We can use set_index to assign key as the index which will allow for easy alignment then groupby to create a dictionary of DataFrames which can be concat:

d = dict(tuple(df.set_index('key', drop=False).groupby('variable')))
out = pd.concat(d, axis=1)

out:

           a                     b                            c                    
    variable key  i  j  k variable  key    i    j    k variable  key    i    j    k
key                                                                                
1          a   1  0  0  0      NaN  NaN  NaN  NaN  NaN        c  1.0  0.0  0.0  0.0
2          a   2  0  0  0        b  2.0  0.0  0.0  0.0      NaN  NaN  NaN  NaN  NaN
3          a   3  0  0  0      NaN  NaN  NaN  NaN  NaN        c  3.0  0.0  0.0  0.0

With a little bit of cleanup we can match the expected output:

  1. drop to get rid of the variable column (since this data has been moved to the header
  2. rename_axis to make index name keys_master instead of keys
  3. swaplevel to make collapsing the MultiIndex easier
  4. reset_index to restore the RangeIndex and make keys_master a column
  5. Index.map to collapse the MultiIndex into a single level
out = (
    pd.concat(
        dict(tuple(df.set_index('key', drop=False).groupby('variable'))),
        axis=1
    ).drop('variable', axis=1, level=1)
        .rename_axis(index='keys_master')
        .swaplevel(axis=1)
        .reset_index()
)
out.columns = out.columns.map(''.join)

out:

   keys_master  keya  ia  ja  ka  keyb   ib   jb   kb  keyc   ic   jc   kc
0            1     1   0   0   0   NaN  NaN  NaN  NaN   1.0  0.0  0.0  0.0
1            2     2   0   0   0   2.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN
2            3     3   0   0   0   NaN  NaN  NaN  NaN   3.0  0.0  0.0  0.0

However, it can be much more simply achieved with pivot and sort_index assuming key variable relationship is unqiue:

out = (
    # Duplicate key column
    df.assign(keys_master=df['key'])
        # Pivot to make variable the columns aligned on keys_master
        .pivot(index='keys_master', columns='variable')
        # Sort index level 1 so variable is grouped together
        .sort_index(level=1, axis=1, sort_remaining=False)
        # Restore RangeIndex and make keys_master a column again
        .reset_index()
)
# Collapse MultiIndex
out.columns = out.columns.map(''.join)

out:

   keys_master  keya   ia   ja   ka  keyb   ib   jb   kb  keyc   ic   jc   kc
0            1   1.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN   1.0  0.0  0.0  0.0
1            2   2.0  0.0  0.0  0.0   2.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN
2            3   3.0  0.0  0.0  0.0   NaN  NaN  NaN  NaN   3.0  0.0  0.0  0.0
user_154978880
Related