Reducing Code Duplication: Loops and Lists

Viewed 26

What would be the best way to reduce the code duplication in the following example. I could only reduce 1st two parts, but I wonder how to make the whole function elegant.

def math_task(data):
    answer = []
    # raise to the third power 
    for elem in data:
        answer += [elem ** 3]
    # take the remainder of the division
    for i in range(len(answer)):
        answer[i] = answer[i] % 5
    # add the original list to the remainder
    for i in range(len(answer)):
        answer[i] = answer[i] + data[i]
    # return the result
    return answer

math_task(test_data)
# print(math_task([1, 4, 5, 9]))
2 Answers

refining the comprehension approach, no index is needed:

def math_task(data):
    return [((elem **3) % 5) + elem for elem in data]

I think you could simply do it like this. You don't need to do a loop for every arithmetic operation you do. You can do all of them in one loop like so:

def math_task(data):
    answers = []
    for i in range(len(data)):
        answers.append(data[i] + (data[i]**3) % 5)
    return answers

However I think the most elegant or "pythonic" way to do this would be by using list comprehension like so.

def math_task(data):
    return [data[i] + (data[i]**3) % 5 for i in range(len(data))]
Related