Search code examples
pythonpandasdataframenumpyaverage

Take average for A,B,C,D,E,F and add a column based on average in python


Take average for A,B,C,D,E,F and add a column based on average.

x=[A,B,C,D,E,F]
def avgfun(x):
    if (np.mean(x,axis=0)) >= 4.5:
        return "High"
    elif (np.mean(x,axis=0)>=3.5) & (np.mean(x,axis=0)<4.5):
        return "Moderate"
    elif (np.mean(x,axis=0)>=2.5) & (np.mean(x,axis=0)<3.5):
        return "Passive"
    else:
        return "Low"

df["Average"]= df[x].mean(axis=1)

Solution

  • It is easier to calculate the mean at first and then apply a function to get new values.

    df = pd.DataFrame({'A': np.random.randint(0, 10, 100),
                       'B': np.random.randint(0, 10, 100),
                       'C': np.random.randint(0, 10, 100),
                       'D': np.random.randint(0, 10, 100),
                       'E': np.random.randint(0, 10, 100),
                       'F': np.random.randint(0, 10, 100)})
    
    def avgfun(x):
        if x >= 4.5:
            return "High"
        elif (x >= 3.5) & (x < 4.5):
            return "Moderate"
        elif (x >= 2.5) & (x < 3.5):
            return "Passive"
        else:
            return "Low"
    l=['A', 'B', 'C', 'D', 'E', 'F']
    df['average'] = df[l].mean(1).apply(lambda x: avgfun(x))