python - 如何编写numpy函数，在python中将它的第k对角线返回为0？

Question

NumPy 函数，它接受一个矩阵和一个整数 k 作为参数，并返回该矩阵的一个副本，使得其第 k 个对角线下的所有元素都为零。主对角线是第 0 条，它上面的对角线编号为 1、2、...。. .，下面的对角线编号为 -1、-2、. . .. 也就是说，1 是主对角线上方的第一条对角线，2 是主对角线上方的第二条对角线，以此类推。-1 是主对角线下方的第一条对角线，-2 是主对角线下方的第二条对角线。

例子：

我想我想出了如何做对角线，但后来我被卡住了：

def EX4(N):
Z=np.zeros((N,N))  #create a array N x N which contains only zero elements


#Set the diagonal to 1
for i in range(k,N):
     Z[i,i]=1 
return Z

k=1 EX4(5)

标签： pythonnumpy

Answer 1

您可能正在寻找numpy.triu()，它完全支持这种行为。

>>> import numpy as np
>>> x = np.array([[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1]])
>>> np.triu(x, 0)
array([[1, 1, 1, 1, 1],
       [0, 1, 1, 1, 1],
       [0, 0, 1, 1, 1],
       [0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1]])
>>> np.triu(x, 1)
array([[0, 1, 1, 1, 1],
       [0, 0, 1, 1, 1],
       [0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0]])
>>> np.triu(x, 2)
array([[0, 0, 1, 1, 1],
       [0, 0, 0, 1, 1],
       [0, 0, 0, 0, 1],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]])