Pandas：按列中的观察数量扩展DataFrame

Question

Stata has the function expand which adds rows to a database corresponding to values in a particular column. For example:

I have:

df = pd.DataFrame({"A":[1, 2, 3], 
                   "B":[3,4,5]})

   A  B
0  1  3
1  2  4
2  3  5

What I need:

df2 = pd.DataFrame({"A":[1, 2, 3, 2, 3, 3], 
                    "B":[3,4,5, 4, 5, 5]})

   A  B
0  1  3
1  2  4
2  3  5
3  2  4
4  3  5
6  3  5

The value in df.loc[0,'A'] is 1, so no additional row is added to the end of the DataFrame, since B=3 is only supposed to occur once.

The value in df.loc[1,'A'] is 2, so one observation was added to the end of the DataFrame, bringing the total occurrences of B=4 to 2.

The value in df.loc[2,'A'] is 3, so two observations were added to the end of the DataFrame, bringing the total occurrences of B=5 to 3.

I've scoured prior questions for something to get me started, but no luck. Any help is appreciated.

Answer 1

There are a number of possibilities, all built around np.repeat :

def using_reindex(df):
    return df.reindex(np.repeat(df.index, df['A'])).reset_index(drop=True)

def using_dictcomp(df):
    return  pd.DataFrame({col:np.repeat(df[col].values, df['A'], axis=0) 
                          for col in df})

def using_df_values(df):
    return pd.DataFrame(np.repeat(df.values, df['A'], axis=0), columns=df.columns)

def using_loc(df):
    return df.loc[np.repeat(df.index.values, df['A'])].reset_index(drop=True)

---

For example,

In [219]: df = pd.DataFrame({"A":[1, 2, 3], "B":[3,4,5]})
In [220]: df.reindex(np.repeat(df.index, df['A'])).reset_index(drop=True)
Out[220]: 
   A  B
0  1  3
1  2  4
2  2  4
3  3  5
4  3  5
5  3  5

---

Here is a benchmark on a 1000-row DataFrame; the result being a roughly 500K-row DataFrame:

In [208]: df = make_dataframe(1000)

In [210]: %timeit using_dictcomp(df)
10 loops, best of 3: 23.6 ms per loop

In [218]: %timeit using_reindex(df)
10 loops, best of 3: 35.8 ms per loop

In [211]: %timeit using_df_values(df)
10 loops, best of 3: 31.3 ms per loop

In [212]: %timeit using_loc(df)
1 loop, best of 3: 275 ms per loop

This is the code I used to generate df :

import numpy as np
import pandas as pd

def make_dataframe(nrows=100):
    df = pd.DataFrame(
        {'A': np.arange(nrows),
         'float': np.random.randn(nrows),
         'str': np.random.choice('Lorem ipsum dolor sit'.split(), size=nrows),
         'datetime64': pd.date_range('20000101', periods=nrows)},
        index=pd.date_range('20000101', periods=nrows))
    return df

df = make_dataframe(1000)

---

If there are only a few columns, using_dictcomp is the fastest. But note that using_dictcomp assumes df has unique column names. The dict comprehension in using_dictcomp won't repeat duplicated column names. The other alternatives will work with repeated column names, however.

Both using_reindex and using_loc assume df has a unique index.

---

using_reindex came from cᴏʟᴅsᴘᴇᴇᴅ's using_loc , in an (unfortunately) now deleted post. cᴏʟᴅsᴘᴇᴇᴅ showed it wasn't necessary to manually repeat all the values -- you only need to repeat the index and then let df.loc (or df.reindex ) repeat all the rows for you. It also avoids accessing df.values which could generate an intermediate NumPy array of object dtype if df contains columns of multiple dtypes.