python - Pandas:通过列中的观察数扩展 DataFrame

Question

Stata 具有函数 expand它将行添加到与特定列中的值相对应的数据库。例如:

我有:

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

   A  B
0  1  3
1  2  4
2  3  5

我需要什么:

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

df.loc[0,'A'] 中的值为 1，因此没有额外的行添加到 DataFrame 的末尾，因为 B=3 应该只出现一次。

df.loc[1,'A'] 中的值为 2，因此在 DataFrame 的末尾添加了一个观察值，使 B=4 的总出现次数变为 2。

df.loc[2,'A'] 中的值为 3，因此在 DataFrame 的末尾添加了两个观察值，使 B=5 的总出现次数变为 3。

我已经搜索了之前的问题以寻找可以让我开始的东西，但没有运气。感谢您的帮助。

Answer 1

有很多可能性，都是围绕 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)

---

例如，

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

---

这是 1000 行的基准数据框；结果是大约 500K 行的 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

这是我用来生成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)

---

如果只有几列，using_dictcomp 是最快的。但请注意，using_dictcomp 假定 df 具有唯一的列名。 using_dictcomp 中的字典理解不会重复重复的列名。但是，其他替代方案将适用于重复的列名。

using_reindex 和using_loc 都假设df 有一个唯一的索引。

---

using_reindex 来自 cᴏʟᴅsᴘᴇᴇᴅ 的 using_loc，现在(不幸的是)删除帖子。 cᴏʟᴅsᴘᴇᴇᴅ表明没有必要手动重复所有值——你只需要重复索引然后让df.loc(或df.reindex)重复所有行给你。它还避免访问 df.values，如果 df 包含多个数据类型的列，它可能生成 object 数据类型的中间 NumPy 数组。