Change every n-th element of a row in a 2d numpy array depending on the row number

Question

I have a 2d array:

H = 12
a = np.ones([H, H])
print(a.astype(int))

[[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 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 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 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 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 1 1 1 1 1 1 1]
 [1 1 1 1 1 1 1 1 1 1 1 1]]

The goal is, for every row r to substitute every r+1-th (starting with 0th) element of that row with 0. Namely, for the 0th row substitute every 'first' (i.e. all of them) element with 0. For the 1st row substitute every 2nd element with 0. And so on.

It can trivially be done in a loop (the printed array is the desired output):

for i in np.arange(H):
    a[i, ::i+1] = 0

print(a.astype(int))

[[0 0 0 0 0 0 0 0 0 0 0 0]
 [0 1 0 1 0 1 0 1 0 1 0 1]
 [0 1 1 0 1 1 0 1 1 0 1 1]
 [0 1 1 1 0 1 1 1 0 1 1 1]
 [0 1 1 1 1 0 1 1 1 1 0 1]
 [0 1 1 1 1 1 0 1 1 1 1 1]
 [0 1 1 1 1 1 1 0 1 1 1 1]
 [0 1 1 1 1 1 1 1 0 1 1 1]
 [0 1 1 1 1 1 1 1 1 0 1 1]
 [0 1 1 1 1 1 1 1 1 1 0 1]
 [0 1 1 1 1 1 1 1 1 1 1 0]
 [0 1 1 1 1 1 1 1 1 1 1 1]]

Can I make use the vectorisation power of numpy here and avoid looping? Or it is not possible?

Answer 1

You can use a np.arange and broadcast modulo over itself

import numpy as np

H = 12

a = np.arange(H)
((a % (a+1)[:, None]) != 0).astype('int')

Output

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