[Tutor] Find Elements in List That Equal A Specific Value

[Steven D'Aprano]

On Thu, 13 May 2010 03:58:45 am Su Chu wrote:

> My problem is as follows:
> I have three lists, one with unique values (list 1), one a sequence
> of values that are not necessarily unique (list2), and a shorter list
> with the unique values of list 2 (list 3). List 1 and List 2 are of
> equal lengths.
>
>
> An example:
> list1 = [ 1, 2, 3, 4, 5, 6 ]
> list2 = [ 2, 2, 2, 5, 6, 6 ]
> list3 = [2, 5, 6]
>
> What I would like to do is find and sum the elements of list 1 given
> its corresponding element in list 2 is equal to some element in list
> 3.

Rather than trying to find a single magic command that happens to do 
exactly what you want, let's break that up into individual steps and 
walk through it by hand.

Start at the end -- you want to match elements from list2 which equals a 
particular value. Where that value comes from doesn't matter.

results = []
for el in list2:
    if el = value:
        results.append(el)


Unfortunately, doing this loses the information we really need: the 
*position* of the element. We don't actually care about the element 
itself. How can we get the position? That's what the enumerate() 
function is for, it takes a list and lazily returns pairs of (position, 
element). So let's adapt the for-loop to do the job:

positions = []
for (pos, el) in enumerate(list2):
    if el == value:
        positions.append(pos)


This can be written as a "list comprehension", which is syntactic sugar 
for a loop:

positions = [pos for (pos, el) in enumerate(list2) if el == value]


Now we have a list of positions. We need to extract the elements in 
list1 at those positions. Here's a for-loop to do it:

elements = []
for pos in positions:  # positions defined above
    elements.append(list1[pos])


And as list comps:

positions = [pos for (pos, el) in enumerate(list2) if el == value]
elements = [list1[pos] for pos in positions]


But we don't need two separate loops, we can combine them into one loop:

elements = []
for (pos, el) in enumerate(list2):
    if el == value:
        elements.append(list1[pos])

And as a single list comp:

elements = [list1[pos] for (pos,el) in enumerate(list2) if el in list3]

And now sum them:

sum(elements)


Now let's put this into a function, so you can call it with different 
values:


def extract_and_sum(list1, list2, value):
    elements = [list1[i] for (i,x) in enumerate(list2) if x in list3]
    return sum(elements)


And now call it in a loop:


sums = []
for value in list3:
    sums.append(extract_and_sum(list1, list2, value))


And you are done.



-- 
Steven D'Aprano