Split/Expand Dataframe based on column values

user_4449593 0
Viewed 640

I have a DataFrame like the below one, with identifiers as a column on top of an existing dateindex.

pd.DataFrame(index = [pd.to_datetime('2021-01-01'), pd.to_datetime('2021-01-01'),pd.to_datetime('2021-01-02'),pd.to_datetime('2021-01-02'), pd.to_datetime('2021-01-03'),pd.to_datetime('2021-01-03')], columns=['id','A', 'B'], data=[['foo',1,5],['bar',8,12],['foo',7,1], ['bar',5,1], ['foo',4,3],['bar',7,1]])

Out[6]: 
             id  A   B
2021-01-01  foo  1   5
2021-01-01  bar  8  12
2021-01-02  foo  7   1
2021-01-02  bar  5   1
2021-01-03  foo  4   3
2021-01-03  bar  7   1

My goal is to create a new sub-dataframes for each of the columns (A and B) except id, with dateIndex as single Index, and id (foo, bar) as column names. The expected output is shown below:

A
Out[9]: 
            foo  bar
2021-01-01    1    8
2021-01-02    7    5
2021-01-03    4    7

B
Out[11]: 
            foo  bar
2021-01-01    5   12
2021-01-02    1    1
2021-01-03    3    1
pandas python
5 Answers
A, B = map(df.set_index('id', append=True).unstack().get, ['A', 'B'])

print(A)

id          bar  foo
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4

print(B)

id          bar  foo
2021-01-01   12    5
2021-01-02    1    1
2021-01-03    1    3
user_23366540

This just simply:

out = df.set_index('id',append=True).unstack('id')
# if you have columns other than `A`,`B`:
# out = df.set_index('id',append=True)[['A','B']].unstack('id')

then you can do

out['A']

which gives:

id          bar  foo
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4

and similarly for out['B']. I found this is much easier and less error prone than hard-coding the variables to A,B.

user_42384080

EDIT: Incorporating the nice idea from @piRSquared to use map in addition to pivot

IIn [58]: A, B = map(lambda column: df[['id', column]].pivot(columns='id', values=column), ['A', 'B'])

In [59]: A
Out[59]: 
id          bar  foo
date                
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4

In [60]: B
Out[60]: 
id          bar  foo
date                
2021-01-01   12    5
2021-01-02    1    1
2021-01-03    1    3
user_65646200

Try with

out = df.set_index('id',append=True).stack().unstack('id').swaplevel(0,1)
A = out.loc['A',:]
A
Out[325]: 
id          bar  foo
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4

Or

d = {x : df[[x,'id']].pivot(columns='id',values=x)  for x in ['A','B']}
d['A']
Out[336]: 
id          bar  foo
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4

Or

df.pivot(columns='id').loc[:,'A']
Out[340]: 
id          bar  foo
2021-01-01    8    1
2021-01-02    5    7
2021-01-03    7    4
user_79645270

You can use pandas.DataFrame.xs to get the values of id foo and bar after set_index id and swaplevel:

>>> A, B = map(df.set_index('id', append=True).swaplevel(0,1).xs, ['foo', 'bar'])
>>> A
            A  B
2021-01-01  1  5
2021-01-02  7  1
2021-01-03  4  3

>>> B
            A   B
2021-01-01  8  12
2021-01-02  5   1
2021-01-03  7   1

Or using level argument in xs that would save swaplevel:

>>> A, B = (df.set_index('id', append=True).xs(ID, level=1) for ID in ['foo', 'bar'])

# This can be made more readable by creating a `partial` function:

>>> from functools import partial
>>> def get_by_ID(df, level, col='id'):
...     func = partial(df.set_index(col, append=True).xs, level=level)
...     return func
>>> A, B = map(get_by_ID(df=df, level=1), ['foo', 'bar'])

Or, simply:

>>> A, B = (df.loc[df.id == ID, ['A', 'B']] for ID in df.id.unique())

Performance

>>> %timeit A, B = map(df.set_index('id', append=True).unstack().get, ['A', 'B'])
2.22 ms ± 33.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

>>> %timeit A, B =  map(df.set_index('id', append=True).swaplevel(0,1).xs, ['foo', 'bar'])
1.88 ms ± 134 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

>>> %timeit A, B = map(get_by_ID(df=df, level=1), ['foo', 'bar'])
1.73 ms ± 54.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

>>> %timeit A, B = (df.loc[df.id == ID, ['A', 'B']] for ID in df.id.unique())
1.69 ms ± 43 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
user_54317910
Related