Add a value to each "other" field ((i + j)% 2 == 0) from the numpy array

Question

I have a numpy m -by- n array, and I would like to add 1.0 to all the [i, j] entries when (i + j) % 2 == 0 , i.e. "for every other square."

I could, of course, just sort out the fields

 import numpy as np a = np.random.rand(5, 4) for i in range(a.shape[0]): for j in range(a.shape[1]): if (i + j) % 2 == 0: a[i, j] += 1.0 

but of course it is very slow.

Any idea on how to improve this?

Answer 1

You can easily perform the operation in two steps, for example

 import numpy as np a = np.zeros((5, 14)) # Even rows, odd columns a[::2, 1::2] += 1 # Odd rows, even columns a[1::2, ::2] += 1 print a 

Answer 2

Here is one way: NumPy broadcasting -

 a[(np.arange(a.shape[0])[:,None] + np.arange(a.shape[1]))%2==0] += 1 

Explanation: We basically create two arrays that are equivalent to the i-th and j-th iterators. Name them I and J

 I = np.arange(a.shape[0]) J = np.arange(a.shape[1]) 

Now, in order to perform the operation between all possible I and J , we create an extension I to 2D by clicking its elements on the first axis and thereby creating a singleton size along its second axis.

Figuratively, the broadcast effect could be done like this:

  I[:,None] : M , 1 J : 1 , N I[:,None] + J : M, N 

So the final setup will be -

 a[(I[:,None] + J)%2==0] += 1 

In other words, avoid comparing with 0 and directly use mod-2 , which will be essentially 0 or 1 -

 a += (np.arange(a.shape[0])[:,None]-1 + np.arange(a.shape[1]))%2 

---

You can also use np.ix_ to handle odd and then even install strings, for example:

 a[np.ix_(np.arange(0,a.shape[0],2),np.arange(0,a.shape[1],2))] += 1 a[np.ix_(np.arange(1,a.shape[0],2),np.arange(1,a.shape[1],2))] += 1 

Answer 3

You can create a mask for "each other" element and apply the addition to the mask.

 # Create mask m00 = np.zeros(a.shape[0], dtype=bool) m00[0::2] = True m01 = np.zeros(a.shape[1], dtype=bool) m01[0::2] = True m0 = np.logical_and.outer(m00, m01) m10 = np.zeros(a.shape[0], dtype=bool) m10[1::2] = True m11 = np.zeros(a.shape[1], dtype=bool) m11[1::2] = True m1 = np.logical_and.outer(m10, m11) m = np.logical_or(m0, m1) a[m] += 1