简体   繁体   中英

How can I reduce the number of if statements in my 2D numpy code?

I have a numpy ndarray ( M ) holding a 2D symmetric matrix that contains N x N numbers (N is an integer), for example

[[1. 2. 0. 0. 2. 2. 2.]
 [2. 1. 2. 0. 0. 2. 2.]
 [0. 2. 1. 2. 0. 0. 2.]
 [0. 0. 2. 1. 2. 0. 0.]
 [2. 0. 0. 2. 1. 2. 0.]
 [2. 2. 0. 0. 2. 1. 2.]
 [2. 2. 2. 0. 0. 2. 1.]]

I want to examine that: set item i,j in M to 0 where j in i+1, i+2, i+3 to zero in M

I have this code, that do this:

for i in range(len(M)):
     for j in range(len(M)):
        if M[i, j] != 0:
            if j == len(M) - 3:
                M[i, j+2] = 0
                M[i, j+1] = 0
            elif j == len(M) - 2:
                M[i, j+1] = 0
            elif j == len(M) - 1:
                   continue
            else:
                M[i, j + 3] = 0
                M[i, j + 2] = 0
                M[i, j + 1] = 0
        else:
            continue

It works pretty well, but I want to reduce my if statements. I read about M[i,j:j+3] , but I don't really know how to use it.

I'm waiting the following output matrix:

[[1. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]
 [0. 2. 0. 0. 0. 0. 2.]
 [0. 0. 2. 0. 0. 0. 0.]
 [2. 0. 0. 0. 1. 0. 0.]
 [2. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]]

This should work:

for i in range(len(M)):
     for j in range(len(M)):
        if M[i, j] != 0:
            M[i, (j+1):(j+4)] = 0
print (M)

Output:

[[1. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]
 [0. 2. 0. 0. 0. 0. 2.]
 [0. 0. 2. 0. 0. 0. 0.]
 [2. 0. 0. 0. 1. 0. 0.]
 [2. 0. 0. 0. 2. 0. 0.]
 [2. 0. 0. 0. 0. 2. 0.]]

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM