pythonarraysreturn

Squaring all elements in a list


I am told to

Write a function, square(a), that takes an array, a, of numbers and returns an array containing each of the values of a squared.

At first, I had

def square(a):
    for i in a: print i**2

But this does not work since I'm printing, and not returning like I was asked. So I tried

    def square(a):
        for i in a: return i**2

But this only squares the last number of my array. How can I get it to square the whole list?


Solution

  • You could use a list comprehension:

    def square(list):
        return [i ** 2 for i in list]
    

    Or you could map it:

    def square(list):
        return map(lambda x: x ** 2, list)
    

    Or you could use a generator. It won't return a list, but you can still iterate through it, and since you don't have to allocate an entire new list, it is possibly more space-efficient than the other options:

    def square(list):
        for i in list:
            yield i ** 2
    

    Or you can do the boring old for-loop, though this is not as idiomatic as some Python programmers would prefer:

    def square(list):
        ret = []
        for i in list:
            ret.append(i ** 2)
        return ret