How to change diagonal values when row sum is a certain number-CodePudding

I believe this is a simple question for people with some level of experience in Python. I have a dataframe like this. What I want is that, if the row sums to 0, the diagonal value changes to 1. E.g. df.loc[2,2] and df.loc[3,3] become 1.

I had the code like this, but I know the if line must have several issues.

for i, j in zip(range(len(df.index)), range(len(df.columns))):
    if i == j & df.iloc[[i]].sum(axis = 1) == 0:
        df.loc[i, j] = 1

CodePudding user response：

Assuming you have a square dataframe, we can create a boolean mask which represent the condition where rows sum to 0, then use this mask to update the diagonal values of the dataframe:

m = df.sum(1).eq(0)
df.values[m, m] = 1

          1    2    3         4         5         6
1  0.997476  0.0  0.0  0.002414  0.000110  0.000000
2  0.000000  1.0  0.0  0.000000  0.000000  0.000000
3  0.000000  0.0  1.0  0.000000  0.000000  0.000000
4  0.000000  0.0  0.0  0.997281  0.002521  0.000173
5  0.000000  0.0  0.0  0.000693  0.999222  0.000054
6  0.000046  0.0  0.0  0.035641  0.964293  0.000000