Using For Loop and If Statement Together

Question

I am trying to get a function, "mock_function", to iterate through the array, "x", and either square or cube its elements, depending on if the element is greater than 5 or not. I need the output to be an array as well, containing the elements squared or cubed. Instead I am getting a single number, so the iteration is clearly not working. Can someone kindly show me what the code should actually be?

def mock_function(x):
    empty_list=[]
    for i in x:
        if i>5:
            mock_answer=i**2
        else:
            mock_answer=i**3
            return mock_answer
    return empty_list.append(mock_answer)

x=np.array([7,4,1,10,3,6])

mock_function(x)

Asked By: TempAcc

||

Source

Answer 1

Append to the list on each iteration instead of returning. Only return the list at the end of the function, after the loop.

def mock_function(x):
    res = []
    for i in x:
        if i>5:
            mock_answer=i**2
        else:
            mock_answer=i**3
        res.append(mock_answer)
    return res

This can also be written more concisely with a list comprehension:

def mock_function(x):
    return [i ** (2 if i > 5 else 3) for i in x]
# i ** (3 - (i > 5)) also works

Answered By: Unmitigated

Answer 2

Using List Comprehension

temp = [22, 4, 45, 50, 3, 69, 4, 44, 1]
[x**2 if x >= 5 else x**3 for x in temp]

Output:

[484, 64, 2025, 2500, 27, 4761, 64, 1936, 1]

Answered By: S M

Using For Loop and If Statement Together

Question:

Answers: