Zero out matrix higher diagonal using numpy

Question

Let's say we have a n x n matrix. Taking n=4 as example:

x x x x
x x x x
x x x x
x x x x

This is what I want to achieve:

When cut=1, given from function parameter, the matrix becomes:

x x x x
x x x x
x x x x
x x x 0

When cut=3:

x x x x
x x x 0
x x 0 0
x 0 0 0

When cut=5:

x x 0 0
x 0 0 0
0 0 0 0
0 0 0 0

As we can see, the diagonal is being cut like a forward-slash, everything under the first slash would be zeros out.

I am using numpy's matrix facility to generate matrices, but I don't know how to code such algorithm. Any help please?

You can always assume this matrix will be an n x n matrix and cut < 2n - 1

Answer 1

You can use np.tri to generate a matrix that contains ones below a given diagonal and zeros above. Since you want the bottom-right corner zeroed out, we have to flip left and right:

bottom_right = lambda N, k: np.fliplr(np.tri(N, k=k-N)) == 1

For example, bottom_right(4, 2) creates the following boolean matrix:

array([[False, False, False, False],
       [False, False, False, False],
       [False, False, False,  True],
       [False, False,  True,  True]], dtype=bool)

You can use that as a slice to zero out the corner:

a = np.ones((4, 4))
a[bottom_right(4, 2)] = 0

a is now

array([[ 1.,  1.,  1.,  1.],
       [ 1.,  1.,  1.,  1.],
       [ 1.,  1.,  1.,  0.],
       [ 1.,  1.,  0.,  0.]])

Another way:

A one-liner based on Mathias711's answer:

f = lambda a, k: np.fliplr(np.triu(np.fliplr(a), k=k-a.shape[0]+1))