Skip to content
Advertisement

Python: for loop that drops a column to meet condition

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?

Advertisement

Answer

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)
Advertisement