Learning Scientific Programming with Python (2nd edition)

E6.1: A "comb" Function

To create a "comb'' of values in an array of length $N$ for which every $n$th element is one but with zeros everywhere else:

import numpy as np
N, n = 101, 5
def f(i):
    return (i % n == 0) * 1

comb = np.fromfunction(f, (N,), dtype=int)
print(comb)

The output is:

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