Is there an easy way to expand/complete a pandas DataFrame to include missing observations with multiple columns?

user_8715278 0
Viewed 655

I have a DataFrame that looks like this:

>>> df = pd.DataFrame({
    'category1': list('AABAAB'),
    'category2': list('xyxxyx'),
    'year': [2000, 2000, 2000, 2002, 2002, 2002],
    'value': [0, 1, 0, 4, 3, 4]
})

>>> df
  category1 category2  year  value
0         A         x  2000      0
1         A         y  2000      1
2         B         x  2000      0
3         A         x  2002      4
4         A         y  2002      3
5         B         x  2002      4

I'd like to expand the data to include missing years in a range. For example, if the range were range(2000, 2003), the expanded DataFrame should look like this:

  category1 category2  year  value
0         A         x  2000    0.0
1         A         y  2000    1.0
2         B         x  2000    0.0
3         A         x  2001    NaN
4         A         y  2001    NaN
5         B         x  2001    NaN
6         A         x  2002    4.0
7         A         y  2002    3.0
8         B         x  2002    4.0

I have tried an approach using pd.MultiIndex.from_product, but that creates rows that aren't valid combinations of category1 and category2 (for example, B and y shouldn't go together). Using from_product and then filtering is too slow for my actual data, which includes many more combinations.

Is there an easier solution to this that can scale well?

Edit

This is the solution I ended up going with, trying to generalize the problem a bit:

id_cols = ['category1', 'category2']

df_out = (df.pivot_table(index=id_cols, values='value', columns='year')
            .reindex(columns=range(2000, 2003))
            .stack(dropna=False)
            .sort_index(level=-1)
            .reset_index(name='value'))

  category1 category2  year  value
0         A         x  2000    0.0
1         A         y  2000    1.0
2         B         x  2000    0.0
3         A         x  2001    NaN
4         A         y  2001    NaN
5         B         x  2001    NaN
6         A         x  2002    4.0
7         A         y  2002    3.0
8         B         x  2002    4.0
pandas python
2 Answers

Let us do stack and unstack 

dfout=df.set_index(['year','category1','category2']).\
         value.unstack(level=0).\
         reindex(columns=range(2000,2003)).\
         stack(dropna=False).to_frame('value').\
         sort_index(level=2).reset_index()
  category1 category2  year  value
0         A         x  2000    0.0
1         A         y  2000    1.0
2         B         x  2000    0.0
3         A         x  2001    NaN
4         A         y  2001    NaN
5         B         x  2001    NaN
6         A         x  2002    4.0
7         A         y  2002    3.0
8         B         x  2002    4.0
user_79645270 
#create a separate dataframe and merge with df to get your new form
a = ('A','x',range(2000,2003))
b = ('A','y',range(2000,2003))
c = ('B','x',range(2000,2003))

from itertools import product, chain
res = ((product(*ent)) for ent in (a,b,c))
columns = ['category1','category2','year']
fake = pd.DataFrame(chain.from_iterable(res),columns=columns)
fake.merge(df,on=columns,how='left').sort_values('year',ignore_index=True)

category1   category2   year    value
0   A   x   2000    0.0
1   A   y   2000    1.0
2   B   x   2000    0.0
3   A   x   2001    NaN
4   A   y   2001    NaN
5   B   x   2001    NaN
6   A   x   2002    4.0
7   A   y   2002    3.0
8   B   x   2002    4.0

alternatively :

fake = df.drop_duplicates(['category1','category2']).filter(['category1','category2'])

fake.index = [2001]*len(fake)
#merge two indexes on year    
pd.concat((df.set_index('year'),fake)).sort_index()

Update

You can abstract the process by using the complete function from pyjanitor:

# pip install pyjanitor
import janitor

df.complete(
        ("category1", "category2"),
        {"year": [2000, 2001, 2002]},
        sort = True
)

  category1 category2   year    value
0   A           x       2000    0.0
1   A           x       2001    NaN
2   A           x       2002    4.0
3   A           y       2000    1.0
4   A           y       2001    NaN
5   A           y       2002    3.0
6   B           x       2000    0.0
7   B           x       2001    NaN
8   B           x       2002    4.0
user_71757130
Related