CODEWITHSUNDEEP
pythonnumpy
James Shapiro
James Shapiro

Reputation: 6186

Best way to negate every other entry of the diagonal of a matrix in Python

I want to create a numpy matrix that is the identity with every other diagonal entry negated. A few options are:

Ugly:

n = 10
i_star = np.eye(n)
for i in range(n):
    if i % 2 == 1:
        i_star[i,i] *= -1

Slightly better: 

n = 10
i_star = np.eye(n)
i_star = i_star[[np.arange(1,n,2)],[np.arange(1,n,2)]] *= -1

Is there a more elegant solution?

Upvotes: 2

Views: 158

Answers (3)

cs95
cs95

Reputation: 402253

I think so. Since the matrix is square, you can use the same index list for both the rows and columns.

idx = np.arange(1,n,2) 
i_star[idx,idx] = -1

Even shorter, using np.r_:

i_star[np.r_[1:n:2], np.r_[1:n:2]] = -1


print(i_star)
array([[ 1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0., -1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0., -1.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0., -1.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0., -1.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., -1.]])

Upvotes: 3

lakshayg
lakshayg

Reputation: 2173

Try this

from itertools import cycle, islice
import numpy as np

i_star = np.diag(list(islice(cycle([1, -1]), 10)))

Upvotes: 2

Julien
Julien

Reputation: 15071

Simply use np.diag:

np.diag([1,-1]*5)

Upvotes: 3

Related Questions