CODEWITHSUNDEEP
pythonpandasdataframemulti-index
Will H
Will H

Reputation: 459

Reset secondary index in pandas dataframe to start at 1

Suppose I construct a multi-index dataframe like the one show here:

prim_ind=np.array(range(0,1000))
for i in range(0,1000):
    prim_ind[i]=round(i/4)

d = {'prim_ind' :prim_ind,
     'sec_ind' : np.array(range(1,1001)),
     'a' : np.array(range(325,1325)),
     'b' : np.array(range(8318,9318))}

df= pd.DataFrame(d).set_index(['prim_ind','sec_ind'])

enter image description here

The sec_ind runs sequentially from 1 upwards, but I want to reset this second index so that for each of the prim_ind levels the sec_ind always starts at 1. I have been trying to work out if I can use reset index to do this but am failing miserably.

I know i could iterate over the dataframe to get this outcome but that will be a horrible way to do it and there must be a more pythonic way - can anyone help?

Note: the dataframe i'm working with is actually imported from csv, the code above is just to illustrate this question.

Upvotes: 2

Views: 606

Answers (1)

jezrael
jezrael

Reputation: 862396

You can use cumcount for count categories.

df.index = [df.index.get_level_values(0), df.groupby(level=0).cumcount() + 1]

Or better if want also index names is use MultiIndex.from_arrays:

df.index = pd.MultiIndex.from_arrays([df.index.get_level_values(0), 
                                      df.groupby(level=0).cumcount() + 1], 
                                      names=df.index.names)
print (df)
                     a     b
prim_ind sec_ind            
0        1         325  8318
         2         326  8319
         3         327  8320
1        1         328  8321
         2         329  8322
         3         330  8323
2        1         331  8324

So column sec_ind is not necessary, you can use also:

d = {'prim_ind' :prim_ind,
     'a' : np.array(range(325,1325)),
     'b' : np.array(range(8318,9318))}

df = pd.DataFrame(d)
print (df.head(8))
     a     b  prim_ind
0  325  8318         0
1  326  8319         0
2  327  8320         0
3  328  8321         1
4  329  8322         1
5  330  8323         1
6  331  8324         2
7  332  8325         2

df = df.set_index(['prim_ind', df.groupby('prim_ind').cumcount() + 1]) \
       .rename_axis(('first','second'))
print (df.head(8))

                a     b
first second           
0     1       325  8318
      2       326  8319
      3       327  8320
1     1       328  8321
      2       329  8322
      3       330  8323
2     1       331  8324
      2       332  8325

Upvotes: 1

Related Questions