Home > other >  Set diagonal values in a Dataframe
Set diagonal values in a Dataframe

Time:12-17

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.

CodePudding user response:

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
  • Related