Is it possible to call a function (not a built in function) inside a list comprehension in python?

Question

Say we have a list:

X = [1,2,3,4,5,6,7,8]

and we have created a function called add():

def add():
  sum = 0
  result = 0
  for e in X:
    sum = sum + e
    return sum
add()

which runs through a list of numbers X, adding the next element in the list to the previous sum. So for each element X[i], we have:

1
3
6
10
15
21
28
36
45

Now, what if I want to put these results in a list again, by making use of a list comprehension. Is it possible to call a function such as add() within a list comprehension, given that it is possible to apply built in functions inside list comprehensions?

I have tried the following:

L = [add() for e in X]
print L

which gives

[None, None, None, None, None, None, None, None, None]

instead of

[1,3,6,10,15,21,28,36,45]

Why am I getting NoneType values in this list?

Answer 1

You could do this with yield to keep with your original format:

    X = [1,2,3,4,5,6,7,8]

    def add():
      sum = 0
      for e in X:
        sum = sum + e
        yield sum

    L = [value for value in add()]
    print L

Answer 2

Yes, it is possible to call functions inside list comprehensions. Your example is fine - it's the add() function that is to be blamed. What you need is make the add() function receive an argument - the list to sum.

def add(elements):
    sum = 0
    for el in elements:
        sum += el
    return sum

This way, the list comprehension would look like this:

L = [add(X[:i+1]) for i in xrange(len(X))]
[1, 3, 6, 10, 15, 21, 28, 36]

This is equivalent to:

L = [add(X[:1]), add(X[:2]), ..., add(X[:8])]

Which turns out to be a list of prefix sums - the thing you want.

Answer 3

Your approach won't work, because your add() is stateless. You need something that maintains state between different invocations of add() , otherwise add() will always produce the same output.

One solution to accomplish what you want is itertools.accumulate() . See Equivalent of Haskell scanl in python for a discussion.

Other answers here have suggested using a list comprehension that involves a range() . While that would work, it would also be inefficient, as an O(n^2) algorithm that recomputes the cumulative sums from scratch for each entry.

Answer 4

#!/usr/bin/env python
"""
produce a list that adds each item to the previous item
this list [1,2,3,4] will produce this list [1,3,6,10]
"""

def accumulate(my_list, previous = 0):
    for i in my_list:
        previous += i
        yield previous

x = [1,2,3,4,5,6,7,8,9]
new_list = []
new_list = [i for i in accumulate(x)]
print x   
print new_list

produces this:

[1, 2, 3, 4, 5, 6, 7, 8, 9]
[1, 3, 6, 10, 15, 21, 28, 36, 45]

Answer 5

Yes, it is possible to call any function in list comprehension.

Also note that in Python 2.x you cannot have print in list comprehension because print is not a function (in Python 2).

To your example - it could look like this, using an add() function that works with shared state (the s variable):

s = 0   # better not use name 'sum' as it is already a builtin function
def add(i):
    global s
    s += i
    return s
X = [1, 2, 3, 4, 5, 6, 7, 8]
print [add(i) for i in X]
# prints [1, 3, 6, 10, 15, 21, 28, 36]
# but beware the global shared state! for example when called again:
print [add(i) for i in X]
# prints [37, 39, 42, 46, 51, 57, 64, 72]

See other answers how to do it without that shared state so you don't get different results when you forget to set s = 0 .