Python - How to expand a Pandas dataframe's rows to include all combinations of values of the key columns?

user_1087924 0
Viewed 1265

Suppose I have the following df:

df=pd.DataFrame({'g1':['a','b','c'],
                'g2':['x','y','z'],
                'val':[1,2,3]})
df
g1  g2  val
a   x   1
b   y   2
c   z   3

whose 'primary key' (using SQL's jargon here) is g1 and g2.

How do I expand the dataframe such that it has a row for every combination of the existing values in g1 and g2. For the rows that are added, leaving np.nan under the val column is fine. So the result will look like the following:

g1  g2  val
a   x   1.0
a   y   NaN
a   z   NaN
b   x   NaN
b   y   2.0
b   z   NaN
c   x   NaN
c   y   NaN
c   z   3.0

I am currently doing the following, but I am sure it is not the best solution:

df.set_index(['g1','g2']).T.stack().unstack().T.reset_index()

I can also do merges, but when the number of keys is large the code looks cumbersome.

pandas python
2 Answers

Use @senderle's fast cartesian_product for performance:

v = cartesian_product(df.g1, df.g2)
idx = pd.MultiIndex.from_arrays([v[:, 0], v[:, 1]])

df.set_index(['g1', 'g2']).reindex(idx)
     val
a x  1.0
  y  NaN
  z  NaN
b x  NaN
  y  2.0
  z  NaN
c x  NaN
  y  NaN
  z  3.0

%timeit df.set_index(['g1','g2']).T.stack().unstack().T
%%timeit
v = cartesian_product(df.g1, df.g2)
idx = pd.MultiIndex.from_arrays([v[:, 0], v[:, 1]])
df.set_index(['g1', 'g2']).reindex(idx)

14.6 ms ± 840 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.56 ms ± 284 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
user_49090870

Answering my own question. The following method should be much faster than the one in my question, but slower than @coldspeed's method. It can be a good alternative when data types of the levels of the index are not the same. 

df.set_index(['g1', 'g2']).reindex(pd.MultiIndex.from_tuples(itertools.product(df.g1, df.g2)))

A rough speed comparison:

  • coldspeed: 100%
  • this: 135%
  • the one in the question: 270%
user_10879240
Related