Replace upper and lower triangle values in a dataframe with zero, or keep only diagonal values

Question

I have the following DataFrame as a toy example:

a = [5,2,6,8]
b = [2,10,19,16]
c = [3,8,15,17]
d = [3,8,12,20]
df  = pd.DataFrame([a,b,c,d], columns = ['a','b','c','d'])
df

I want to create a new DataFrame df1 that keeps only the diagonal elements and converts upper and lower triangular values to zero.

My final dataset should look like:

    a   b   c   d
0   5   0   0   0
1   0   10  0   0
2   0   0   15  0
3   0   0   0   20

Answer 1

You could use numpy.diag:

df = pd.DataFrame(data=np.diag(np.diag(df)), columns=df.columns)
print(df)

Output

   a   b   c   d
0  5   0   0   0
1  0  10   0   0
2  0   0  15   0
3  0   0   0  20

Answer 2

import pandas as pd

def diag(df):
    res_df = pd.DataFrame(0, index=df.index, columns=df.columns)
    for i in range(min(df.shape)): res_df.iloc[i, i] = df.iloc[i, i]
    return res_df