Reformat Dataframe / Add rows when condition is met

Viewed 54

I'm looking to add dataframe rows and edit a column when a condition is met. I want Column B to be only "1's". If the value is greater than one, then add length of rows equal to the number thats > 1, while keeping ColA sorted by date asc. Example below:

Original DF:

   ColA        ColB
2021-03-09       1
2021-03-09       3
2021-03-10       2
2021-03-10       1
2021-03-10       2
2021-03-11       2

Desired DF

   ColA         ColB
2021-03-09       1
2021-03-09       1
2021-03-09       1
2021-03-09       1
2021-03-10       1
2021-03-10       1
2021-03-10       1
2021-03-10       1
2021-03-10       1
2021-03-11       1
2021-03-11       1

any suggestions are much appreciated!

1 Answers

Assuming you only have positive integers in 'ColB' You can re-create the DataFrame from scratch using np.repeat. The repeat takes care of the duplication, so we can assign ColB = 1.

import pandas as pd
import numpy as np

df = (pd.DataFrame(np.repeat(df.ColA, df.ColB))
        .assign(ColB=1))

Alternatively, if you have a non-duplicated Index, you can repeat that and use loc to get the repitition. Useful when you have more than a single column you want to repeat:

df = (df.loc[df.index.repeat(df.ColB)]
        .assign(ColB=1))

         ColA  ColB
0  2021-03-09     1
1  2021-03-09     1
1  2021-03-09     1
1  2021-03-09     1
2  2021-03-10     1
2  2021-03-10     1
3  2021-03-10     1
4  2021-03-10     1
4  2021-03-10     1
5  2021-03-11     1
5  2021-03-11     1
Related