Search code examples
pythonpython-3.xpandasmulti-index

Adding two values of two columns and assigning the result to a third column in a pandas multi-index DataFrame


I have a Pandas dataframe:

a=[1,1,1,2,2,2,3,3,3]
dic={'A':a}

df=pd.DataFrame(dic)

I apply a multi-index to this df:

index=[(1,'a'),(1,'b'),(1,'c'),(2,'a'),(2,'b'), (2, 'c'),(3,'a'),(3,'b'), (3,'c')]
df.index=pd.MultiIndex.from_tuples(index, names=['X','Y'])

I add a new column:

df['B']='-'

Now I have a df:

       A   B 
X Y          
1 a    1   -
  b    1   -
  c    1   -
2 a    2   -
  b    2   -
  c    2   -
3 a    3   -
  b    3   -
  c    3   -

Essentially, I want to cycle through level='X' of the multi-index, adding one level to another, and then assigning the values to column='B'

Here's how I was thinking about doing it:

dex=[]
for idx, select_df in df.groupby(level=0):
    dex.append(idx)
#gives me a list of level='X' keys

dex_iter=iter(dex)
#creates an iterator from that list

last=next(dex_iter)
#gives me the first value of that list of keys, and moves the iterator to the next value

for i in dex_iter:
    
    df.loc[i,'B']=df.loc[i,'A']+df.loc[last,'A']
    last=i

My EXPECTED result is:

      A   B
X Y        
1 a   1   -
  b   1   -
  c   1   -
2 a   2   3
  b   2   3
  c   2   3
3 a   3   5
  b   3   5
  c   3   5

Instead, what I get is:

      A    B
X Y        
1 a   1    -
  b   1    -
  c   1    -
2 a   2  NaN
  b   2  NaN
  c   2  NaN
3 a   3  NaN
  b   3  NaN
  c   3  NaN

This is obviously due to some peculiarity with assigning the values to the multi-index. But I can't find a way to resolve this issue.


Solution

  • Let's try groupby, first, and shift:

    df.groupby(level=0)['A'].first().shift()
    
    X
    1    NaN
    2    1.0
    3    2.0
    Name: A, dtype: float64
    
    tmp = df.index.get_level_values(0).map(df.groupby(level=0)['A'].first().shift())
    print (tmp)
    # Float64Index([
    #    nan, nan, nan, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0], dtype='float64', name='X')
    

    This gives you the values you need to add to "A" to get "B":

    df['B'] = df['A'] + tmp
    df
    
         A    B
    X Y        
    1 a  1  NaN
      b  1  NaN
      c  1  NaN
    2 a  2  3.0
      b  2  3.0
      c  2  3.0
    3 a  3  5.0
      b  3  5.0
      c  3  5.0