How can I unroll a multi-level pandas dataframe

Viewed 79

I have a pandas dataframe with three columns, the first two columns are factors, and the third column contains counts. I want to 'explode' or 'unroll' the dataframe so that instead of having one line for each unique element of first column, second column, I have the number of rows equal to the sum of the counts column, where each new line has a unique and incrementing identifier number, but I want a separate counter for each level within one of the two columns. Note, this question is similar to How can I 'unroll' a pandas dataframe? which I asked yesterday, but has some additional complications that i failed to recognize the first time, and i'm unable to generalize (for myself) how to extend upon it.

Here are the data frames

data = [['van', 'bc', 1], ['abb', 'bc', 3], ['vic','bc',3], ['cal', 'ab', 1], ['edm', 'ab', 2], ['cal','ab', 2], ['van', 'bc', 1]]
df = pd.DataFrame(data, columns = ['city', 'state', 'count']) 

and I want to turn that into this


data = [['van', 'bc', 'dr0001'], ['abb', 'bc', 'dr0002'], ['abb', 'bc', 'dr0003'], ['abb', 'bc', 'dr0004'],  ['vic', 'bc', 'dr0005'], ['vic', 'bc', 'dr0006'], ['vic', 'bc', 'dr0007'], ['cal', 'ab', 'dr0001'], ['edm', 'ab', 'dr0002'], ['edm', 'ab', 'dr0003'], ['edm', 'ab', 'dr0004'], ['edm', 'ab', 'dr0005'], ['van', 'bc', 'dr0008']]
df = pd.DataFrame(data, columns = ['city', 'state', 'id'])

Thanks

2 Answers

Try this , I think you need an additional groupby with some formatting looking at your output:

out = df.loc[df.index.repeat(df["count"])]
out = (
    out.assign(
        id=out.groupby("state")
        .cumcount()
        .add(1)
        .astype(str)
        .str.zfill(4)
        .radd("dr")
    )
    .drop("count", 1)
    .reset_index(drop=True)
)

print(out)

   city state      id
0   van    bc  dr0001
1   abb    bc  dr0002
2   abb    bc  dr0003
3   abb    bc  dr0004
4   vic    bc  dr0005
5   vic    bc  dr0006
6   vic    bc  dr0007
7   cal    ab  dr0001
8   edm    ab  dr0002
9   edm    ab  dr0003
10  cal    ab  dr0004
11  cal    ab  dr0005
12  van    bc  dr0008
  • generate a list then explode()
  • id increments in state so generate this after DF is right shape
data = [['van', 'bc', 1], ['abb', 'bc', 3], ['vic','bc',3], ['cal', 'ab', 1], ['edm', 'ab', 2], ['cal','ab', 2], ['van', 'bc', 1]]
df = pd.DataFrame(data, columns = ['city', 'state', 'count']) 

# first pass, explode
df2 = (df.assign(id=df["count"].apply(lambda n: [f"dr{i+1:05}" for i in range(n)]))
       .explode("id")
       .drop(columns="count").reset_index(drop=True))

# ids increment within state
df2["id"] = df2.groupby("state")["id"].transform(lambda s: [f"dr{i+1:05}" for i,v in enumerate(s)])

output

city state id
0 van bc dr00001
1 abb bc dr00002
2 abb bc dr00003
3 abb bc dr00004
4 vic bc dr00005
5 vic bc dr00006
6 vic bc dr00007
7 cal ab dr00001
8 edm ab dr00002
9 edm ab dr00003
10 cal ab dr00004
11 cal ab dr00005
12 van bc dr00008
Related