python - Squaring all elements in a list

18
votes

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?

9
pythonarraysreturn
Is this homework? Seems like it is. - Damian Schenkelman
yes it is, I said "I am told to ..." so I thought it was obvious. I took a few attempts on the problem also and could not come up with the format that was asked for so I came here. - user1692517
Please be careful with your use of list and array; those are two different data structures. - Akavall
@Akavall: note that the homework tag is now deprecated and should not be added to questions - David Robinson
@DavidRobinson, Thanks for letting me know. - Akavall

9 Answers

39
votes

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
31
votes

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.

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

Use numpy.

import numpy as np
b = list(np.array(a)**2)
1
votes
def square(a):
    squares = []
    for i in a:
        squares.append(i**2)
    return squares
0
votes

One more map solution:

def square(a):
    return map(pow, a, [2]*len(a))
0
votes
def square(a):
    squares = []
    for i in a:
        squares.append(i**2)
    return squares

so how would i do the square of numbers from 1-20 using the above function

0
votes

you can do

square_list =[i**2 for i in start_list]

which returns

[25, 9, 1, 4, 16]

or, if the list already has values

square_list.extend([i**2 for i in start_list])

which results in a list that looks like:

[25, 9, 1, 4, 16]

Note: you don't want to do

square_list.append([i**2 for i in start_list])

as it literally adds a list to the original list, such as:

[_original_, _list_, _data_, [25, 9, 1, 4, 16]]
-2
votes
n = int(input())

final_list = [n**2 for n in range(1, n+1)]

print(final_list)