Python pandas: Take column of counts and create DataFrame with a row per count

Question

I have a DataFrame with values in columns a and b and a third column with the count of that row. I would like to convert this into a DataFrame (either new or remake the old one) with columns a and b repeated the number of times as is in the count column. It's probably more clear with an example. I have this DataFrame:

import pandas as pd
df = pd.DataFrame({'a' : [1,2,3], 'b' : [0,0,1], 'count' : [3,1,4]})

I am converting it like this:

new_df = pd.DataFrame(columns=df.columns[:-1])
for _, row in df.iterrows():
    num = row['count']
    for i in range(num):
        pd.concat([new_df, row])
        new_df = new_df.append(row[:-1])

This does exactly what I want but seems inelegant to me because of the for-loop inside iterrows. Is there a better or more pythonic way to do this?

Answer 1

You are correct in thinking that iterrows is a very bad sign for Pandas code. Even worse is building up a DataFrame one row at a time like this with pd.concat - the performance implications are dreadful.

Instead of reaching for loops, your first step should be to check if there is a vectorized DataFrame method you could use. In this case.. perhaps there isn't.

Next step is dropping into NumPy. And lo and behold, there is numpy.repeat, which allows you to repeat an array along an axis with another array of counts. We can wrap that up in a function.

def repeat_frame(df, counts):            
    rep_array = np.repeat(df.values, counts, axis=0)
    return pd.DataFrame(rep_array, columns=df.columns)

# would be called with 
repeat_frame(df[['a', 'b']], df['count'])

This runs in 450 µs on my machine compared to your current solution in 20.2 ms for your sample data. On a larger sample size with 1000 rows, it runs in 511 µs compared to 6.86 s for your current solution, roughly 13000 times faster.

---

If you find yourself wanting to use pd.concat or similar in any situation like this to build up a DataFrame row-by-row, stop! There is going to be a faster way.

• Check for vectorized Pandas functions that can help you out
• Ditto for NumPy and SciPy after that

And if you are certain you have some wild iteration logic that cannot be vectorized (which is unlikely), create an empty array of the necessary shape and consider using Numba or Cython to speed up filling it up with a loop.