Python: for loop that drops a column to meet condition-CodePudding

I have a dataframe that looks as follows:

alpha_0  alpha_1  alpha_2  alpha_3
1        2        1        4
2        0        3        8
0        0        0        9

Beta is calculated as ((sum of each row)^2)/10. I want to keep dropping columns until Beta is less than or equal to 1 for all rows.

So far I have

n_alphas=4
for alpha in range(0,n_alphas):
    df.drop(list(values.filter(regex = 'alpha ' str(alpha))), axis = 1, 
    inplace = True)
    
    Beta=(df.sum(axis=1)^2)/10
    print(Beta)

How can I stop the loop when all values of beta are below or equal to 1?

CodePudding user response：

First of all, if you want to calculate the power of a number, do not use ^ operator. It is an XOR Boolean operator. Instead, use ** operator.

This code should work. However, this will also delete last remaining column from the dataframe, if the condition of Beta is not met.

for column in df.columns:
    Beta = (df.sum(axis=1) ** 2) / 10
    if Beta.min() > 1.0:
        df.drop(columns=[column], inplace=True)

If you do not want the last remaining column to be deleted even if the Beta condition is not met, use this code

n_alphas = 4
for alpha in range(0, n_alphas):
    Beta = (df.sum(axis=1) ** 2) / 10
    if Beta.min() > 1.0:
        df.drop(columns=[f"alpha_{alpha}"], inplace=True)

CodePudding user response：

I hope you expecting for this

import pandas as pd

df = pd.DataFrame({'a1': [1,2,0],'a2':[2,0,0],'a3':[1,3,0],'a4':[4,8,9]})
for i in range(len(df)):
    if sum(df.iloc[i])**2 /10 > 1:
        df = df.drop([df.columns[0]], axis=1)