Add difference row in pandas dataframe

Viewed 80

I am wondering if anybody can help me work on my pandas style. I am trying to do an operation on groups of data within a panda, like in the following example.


import pandas as pd

df = pd.DataFrame({
    'idprod': [1, 1, 1, 1, 2, 2, 2],
    'optn': [1, 2, 3, 4, 1, 2, 3],
    'keep': ['a', 'a', 'a', 'a', 6, 6, 6],
    'bla': [7, 6, 5, 4, 0, 9, 8],
    'val': [500, 850, 1100, 1150, 1500, 2500, 4500]})

print(df)


def addtrf(g):
    g.set_index('optn', inplace=True)
    d = g.loc[3, 'val'] - g.loc[1, 'val']
    try: 
         d2 = g.loc[4, 'val'] - g.loc[1, 'val']
    except:
        d2=0
    g.reset_index(inplace=True)

    n = g.loc[g.optn == 3, :]
    n['val'] = d
    n['optn'] = 'diff1'
    del n['bla']
    g = g.append(n)

    n = g.loc[g.optn == 4, :]
    n['val'] = d2
    n['optn'] = 'diff2'
    del n['bla']
    g = g.append(n)

    #print(g)

    return (g)


dfn = df.groupby('idprod').apply(addtrf)
dfn=dfn.drop(columns='idprod')
dfn=dfn.reset_index()
dfn = dfn.drop(columns='level_1')

print(dfn)

So this actually does what it should, however, I don't think it is a nice way to do it. Can you help me to make this more concise? I don't want to use a loop. I have also looked into using .diff() but I may need to do more than just a minus, and I need to keep some columns (keep), and remove others (bla). So similar to https://stackoverflow.com/users/8081945/babz question below this answer: https://stackoverflow.com/a/57378481/13878458

Any help much appreciated! PS: It is very slow on my dataset, and I also have to manually delete the level_1 column which I am getting from the reset_index

0 Answers
Related