pivoting pandas dataframe into prefixed cols, not a MultiIndex

Question

I have a timeseries dataframe that is similar to:

ts = pd.DataFrame([['Jan 2000','WidgetCo',0.5, 2], ['Jan 2000','GadgetCo',0.3, 3], ['Jan 2000','SnazzyCo',0.2, 4],
          ['Feb 2000','WidgetCo',0.4, 2], ['Feb 2000','GadgetCo',0.5, 2.5], ['Feb 2000','SnazzyCo',0.1, 4],
          ], columns=['month','company','share','price'])

Which looks like:

  month   company  share  price
0  Jan 2000  WidgetCo    0.5    2.0
1  Jan 2000  GadgetCo    0.3    3.0
2  Jan 2000  SnazzyCo    0.2    4.0
3  Feb 2000  WidgetCo    0.4    2.0
4  Feb 2000  GadgetCo    0.5    2.5
5  Feb 2000  SnazzyCo    0.1    4.0

I can pivot this table like so:

pd.pivot_table(ts,index='month', columns='company')

Which gets me:

            share                      price                  
company  GadgetCo SnazzyCo WidgetCo GadgetCo SnazzyCo WidgetCo
month                                                         
Feb 2000      0.5      0.1      0.4      2.5        4        2
Jan 2000      0.3      0.2      0.5      3.0        4        2

This is what I want except that I need to collapse the MultiIndex so that the company is used as a prefix for share and price like so:

          WidgetCo_share  WidgetCo_price  GadgetCo_share  GadgetCo_price   ...
month                                                                      
Jan 2000             0.5               2             0.3             3.0   
Feb 2000             0.4               2             0.5             2.5

I came up with this function to do just that but it seems like a poor solution:

def pivot_table_to_flat(df, column, index):
    res = df.set_index(index)
    cols = res.drop(column, axis=1).columns.values
    resulting_cols = []
    for prefix in res[column].unique():
        for col in cols:
            new_col_name = prefix + '_' + col
            res[new_col_name] = res[res[column] == prefix][col]
            resulting_cols.append(new_col_name)

    return res[resulting_cols]

pivot_table_to_flat(ts, index='month', column='company')

What is a better way of accomplishing a pivot resulting in a columns with prefixes as opposed to a MultiIndex?

score 4 · Answer 1 · answered Oct 29 '17 at 09:38

This seems even simpler:

df.columns = [' '.join(col).strip() for col in df.columns.values]

It takes a df with a multiindex column and flattens the column labels, with the df remaining in place.

(ref: @andy-haden Python Pandas - How to flatten a hierarchical index in columns )

score 1 · Accepted Answer · answered Nov 22 '14 at 03:13

I figured it out. Using the data on the MultiIndex makes for a pretty clean solution:

def flatten_multi_index(df):
    mi = df.columns
    suffixes, prefixes = mi.levels
    col_names = [prefixes[i_p] + '_' + suffixes[i_s] for (i_s, i_p) in zip(*mi.labels)]
    df.columns = col_names
    return df

flatten_multi_index(pd.pivot_table(ts,index='month', columns='company'))

The above version only handles a 2D MultiIndex but it could be generalized if needed.

score 1 · Answer 3 · edited Feb 09 '18 at 13:40

An update (as of early 2017 and pandas 0.19.2). You can use .values on a MultiIndex. So, this snippet should flatten MultiIndexs for those in need. The snippet is both too clever but not clever enough: it can handle either the row index or column names from the DataFrame, but it will blow up if the result of getattr(df,way) isn't nested (i.e., a MultiIndex).

def flatten_multi(df, way='index'): # or way='columns'
    assert way in {'index', 'columns'}, "I'm sorry Dave."
    mi = getattr(df, way)
    flat_names = ["_".join(s) for s in mi.values]
    setattr(df, way, flat_names)
    return df

pivoting pandas dataframe into prefixed cols, not a MultiIndex

3 Answers3