pythonpandasdataframediagonal

Set diagonal values in a Dataframe


Similar questions have been asked already but I still can't come up with a solution.

I have a pandas Dataframe of such a shape:

 1     0     0     0     0
 1     0     0     0     0
 1     0     0     0     0
 1     0     0     0     0
-1     0     0     0     0
-1     0     0     0     0
-1     0     0     0     0
-1     0     0     0     0

And I want to make it into one that looks like this:

 1     1     0     0     0
 1     0     1     0     0
 1     0     0     1     0
 1     0     0     0     1
-1     1     0     0     0
-1     0     1     0     0
-1     0     0     1     0
-1     0     0     0     1

I have tried np.fill_diagonal(df.values,1) but that didn't work.


Solution

  • Assuming the first column is indeed a column and that you want to ignore it.

    Use :

    n_ignore = 1
    a = df.to_numpy()
    idx = np.arange(a.shape[0])
    col = np.arange(a.shape[0])%(a.shape[1]-n_ignore)+n_ignore
    a[idx, col] = 1
    
    df2 = pd.DataFrame(a, index=df.index, columns=df.columns)
    

    If you rather want to use fill_diagonal per group:

    def fill_diag(g):
        a = g.to_numpy()
        np.fill_diagonal(a, 1)
        return pd.DataFrame(a, index=g.index, columns=g.columns)
        
    (df.set_index(0)
       .groupby(level=0, group_keys=False).apply(fill_diag)
       .reset_index()
    )
    

    Output:

       0  1  2  3  4
    0  1  1  0  0  0
    1  1  0  1  0  0
    2  1  0  0  1  0
    3  1  0  0  0  1
    4 -1  1  0  0  0
    5 -1  0  1  0  0
    6 -1  0  0  1  0
    7 -1  0  0  0  1