Array in a program

Question

So, first I like to state that I'm very new to Python, so this probably will be an easy question. I'm exercising my programming skills and am trying to write a simple program. It goes like this:

def cal(alpha, k):
    sqrt = np.zeros(len(alpha))
    for i in range(len(alpha)):
        sqrt[i] = np.sqrt(alpha[i])*k
        return sqrt

Which, for a certain alpha and k, should give me an array sqrt[]. For example, if alpha = [1,4,9] and k = 3, the answer should be [3,6,9]. However, when executed in Python, it gives [3,0,0].

My question is: why? I know I can get what I want if I simply put

 def cal(alpha, k):
     sqrt = np.zeros(len(alpha))
     sqrt= np.sqrt(alpha)*k
     return sqrt

but I want to know where my mistake in reasoning is.

Thanks!

Answer 1

Indentation is important in Python. As soon as your function reaches the return statement, it will stop and return an object.

In your for loop, return is met at the end of the first iteration but before the start of the second, which is why you find only the first value filled in [3, 0, 0].

Therefore, simply unindent the last line. This will ensure return is only processed after your for loop is completed:

import numpy as np

def cal(alpha, k):
    sqrt = np.zeros(len(alpha))
    for i in range(len(alpha)):
        sqrt[i] = np.sqrt(alpha[i])*k
    return sqrt

cal([1, 4, 9], 3)

# array([ 3.,  6.,  9.])