问题:

I have two python lists,

A = [ 1, 2, 3, 4, 5 ]
B = [ True, False, False, True, True ]

lists A and B are the same length.

I want to sum only the elements in A that correspond to True elements in B. I know I can do that with something like:

sum([A[x] for x in xrange(len(A)) if B[x]])

but I was wondering if there was a more elegant solution that didn't involve looping over elements in each list?

回答1:

Using itertools.compress:

>>> from itertools import compress
>>> sum(compress(A, B))
10

The implementation of itertools.compress is described on the linked page. It's short and simple, so you don't have to import itertools ¹:

>>> sum(a for a, b in zip(A, B) if b)
10

_{¹ OTOH, itertools.compress is implemented in C, and therefore should be faster}

Generator expression + sum:

>>> A = [1, 2, 3, 4, 5]
>>> B = [True, False, False, True, True]
>>> sum(a for a,b in zip(A, B) if b)
10

Code

A = [ 1, 2, 3, 4, 5 ]
B = [ True, False, False, True, True ]

sum_ = sum([numa for (numa, numb) in zip(A, B) if numb])
print(sum_)

Out

How it works:

If we expand the list comprehension a little more we can see what is happening a little more clearly.

nums_to_sum = []
for numa, numb in zip(A, B):
     if numb:
         nums_to_sum.append(numa)
sum_ = sum(nums_to_sum)