storing output of a for loop that contains multiple columns in a list-CodePudding

I have the following dataframe:

data = {'id':[1,1,2,2,2,3,3,3],
                  'var':[10,12,8,11,13,14,15,16],
                  'ave':[2,2,1,4,3,5,6,8],
                  'age': [10,10,13,13,20,20,20,11],
                  'weight':[2.5,1.1,2.1,2.2,3.5,3.5,4.2,1.3],
                  'con': [3.1,2.3,4.5,5.5,6.5,7.5,4.7,7.1]}

df = pd.DataFrame(data)

by the code below, I want to run a for loop with 6 time iterations and store the results,capita, in a list but I got the erorr KeyError 1. I have tried with dictionary and dataframe but one of them work:

ls = ([])
for i in [1,6]:
    capita = df.groupby('age') \
            .apply(lambda x: x[['con']].mul(x.weight, 0).sum() / (x.weight).sum()) \
             \
            .reset_index()\
                .rename(columns={"con":"ave"})
    df["ave"] =df["age"].map( df.groupby(['age']).
                                   apply(lambda x: np.average(x['con'], weights=x['weight'])))   
    df['con'] =df['var']*df['ave']/df.groupby('id')['ave'].transform('sum')
    ls[i]=capita[i]

CodePudding user response：

The error you are mention is caused because you are trying to access an index of the list ls that does not exist.

Try this:

ls = []
for i in range(6):
    capita = df.groupby('age') \
            .apply(lambda x: x[['con']].mul(x.weight, 0).sum() / (x.weight).sum()) \
             \
            .reset_index()\
                .rename(columns={"con":"ave"})
    df["ave"] =df["age"].map( df.groupby(['age']).
                                   apply(lambda x: np.average(x['con'], weights=x['weight'])))   
    df['con'] =df['var']*df['ave']/df.groupby('id')['ave'].transform('sum')
    ls.append(capita)