Ben Mabey Ben Mabey -4 years ago 165
Python Question

pivoting pandas dataframe into prefixed cols, not a MultiIndex

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
?

Answer Source

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.

Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download