Alf
Alf

Reputation: 1989

Replacing for-loop in python when index is part of a array-manipulating function

I have a multidimensional numpy array which I need to modify such that its elements are modified as a function of the index of one of the dimensions only. I can of course do that with a for loop, as in the following simplified example

import numpy as np
a = np.ones( (2,10) )
for ii in range(a.shape[1]):
    a[:,ii] *= ii

If the array becomes very large, this might slow down the execution and I was wondering if there are some clever ways to avoid using a for loop?

Upvotes: 0

Views: 25

Answers (1)

Mateen Ulhaq
Mateen Ulhaq

Reputation: 27201

Construct another array to hold the scaling factors, then broadcast and multiply:

scale = np.arange(a.shape[1])
a *= scale

Upvotes: 1

Related Questions