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

Question

NumPy function which takes a matrix and an integer k as parameters and returns a copy of this matrix such that all its elements under the kth diagonal are zero. The main diagonal is the 0th one, the diagonal above it are numbered 1, 2, . . ., and the diagonal below are numbered −1, −2, . . .. That is, 1 is the first diagonal above the main diagonal, 2 is the second diagonal above the main diagonal, etc. −1 is the first diagonal below the main diagonal, −2 is the second diagonal below the main diagonal.

example:

I think I figured out how to do a diagonal line, but then I am stuck:

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)

Answer 1

You are probably looking for numpy.triu() , which supports exactly this behavior.

>>> 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]])

Answer 2

You can do something along the lines of the following.

z = np.zeros((4,4))

def f(k):
    for i in range(4):
        for j in range(4):
            if j > i + k:
                z[i,j] = 1

---

Edit: Just seen the other answer. I wasn't aware numpy had this functionality. It would be best to use numpy's methods, however I will leave the answer here still if you were curious for a way to do it yourself.