Squaring all elements in a list

Question

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 1

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

Solution 2

Use a list comprehension (this is the way to go in pure Python):

>>> l = [1, 2, 3, 4]
>>> [i**2 for i in l]
[1, 4, 9, 16]

Or numpy (a well-established module):

>>> numpy.array([1, 2, 3, 4])**2
array([ 1,  4,  9, 16])

In numpy, math operations on arrays are, by default, executed element-wise. That's why you can **2 an entire array there.

Other possible solutions would be map-based, but in this case I'd really go for the list comprehension. It's Pythonic :) and a map-based solution that requires lambdas is slower than LC.

Solution 3

import numpy as np
a = [2 ,3, 4]
np.square(a)

Solution 4

Use numpy.

import numpy as np
b = list(np.array(a)**2)

Solution 5

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

View more solutions

Comments

[Waleed Khan]

Numpy for such a trivial problem seems like overkill.

[Dr. Jan-Philip Gehrcke]

Good that you point out a lot of methods. However, most established solutions are based on list comprehension or numpy. For performance of map in combination with lambda, have a look at stackoverflow.com/questions/1247486/…

[tacaswell]

Fair point, but where you have the need to square lists, you soon need to start doing other operations with them and there is no reason re-invent the wheel.

[user1692517]

Thank you! I used the comprehension method. Will be looking more into that method.

[Benjamin Hodgson]

"This is not as idiomatic as some Python programmers would prefer" - I completely agree, but it's worth pointing out that there are situations in which the only practical option is to append to a list. The best example I can think of is if the generator needs to 'remember' the numbers it has previously returned so as not to return duplicates or get into a cycle.

[user1936752]

Is there a way to avoid having to use np.array and just use array(a)**2?

[wizzwizz4]

@user1936752 No, because the "to the power of" function on standard arrays isn't defined. However, it's defined on numpy arrays, because they aren't really arrays; they are a special data type created by the numpy module.

[AChampion]

Note: square_list.extend() does not return anything as this is an in place operation.

[AChampion]

Also, itertools.repeat(2) would work instead of creating a list of 2s

[9 Guy]

@AChampion you are correct, i'll edit to make it more accurate. Also its funny to look at myself from 4 years ago talk wow

[Coder]

up voted for the O(1) space complexity map solution

[Brandt]

Add some text around your code: either tell us why you added this answer (supposedly, non of the others satisfied you) or make # in your code block since you go for simplicity.

[General Grievance]

This is incorrect. You are supposed to square every element in a given list. This just generates a list of all perfect squares up to n^2. Try it online!