Python Pandas : Balance an unbalanced dataset (for panel analysis)

Question

我知道这可能很容易做到。我可以在 Stata 中完成，但我正在尝试转向 Python。

我有一个不平衡的大数据集。它看起来像这样:

我需要获取如下数据集:

欢迎任何指导。非常感谢!

Answer 1

一种方法是使用 set_index、reindex 使用 pd.MultiIndex.from_product 和 将“年份”设置为另一级索引>reset_index 将“year”中的数据作为列。

具有相同结构的示例数据框:

import pandas as pd

df = pd.DataFrame( {'year':[2003,2004,2002,2004,2005,2006],
                    'city_code':['a']*2+['b']*4,
                    'total_tax':pd.np.random.randint(100,1000,6)},
                   index=pd.Index(data=[9]*2+[54]*4,name='id_inf'))
print(df)
       city_code  total_tax  year
id_inf                           
9              a        417  2003
9              a        950  2004
54             b        801  2002
54             b        218  2004
54             b        886  2005
54             b        855  2006

现在您可以使用以下方法创建df_balanced:

df_balanced = (df.set_index('year',append=True)
                 .reindex(pd.MultiIndex.from_product([df.index.unique(),
                                                      range(df.year.min(),df.year.max()+1)],
                                                     names=['id_inf','year']))
                 .reset_index(level=1))

你会得到:

print (df_balanced)
        year city_code  total_tax
id_inf                           
9       2002       NaN        NaN
9       2003         a      417.0
9       2004         a      950.0
9       2005       NaN        NaN
9       2006       NaN        NaN
54      2002         b      801.0
54      2003       NaN        NaN
54      2004         b      218.0
54      2005         b      886.0
54      2006         b      855.0

要填充NaN，有不同的方法，但这里有两种方法。对于“city_code”列，您可以使用 groupby 和 transform 与 max 来获取值，对于“total_tax”列，只需 fillna 用 0 例如:

df_balanced['city_code'] = df_balanced.groupby(level=0)['city_code'].transform(max)
df_balanced['total_tax'] = df_balanced['total_tax'].fillna(0)

print (df_balanced)
        year city_code  total_tax
id_inf                           
9       2002         a        0.0
9       2003         a      417.0
9       2004         a      950.0
9       2005         a        0.0
9       2006         a        0.0
54      2002         b      801.0
54      2003         b        0.0
54      2004         b      218.0
54      2005         b      886.0
54      2006         b      855.0