2

A not too difficult question, I hope, from a beginner in Python.

I have a main list, listA, and I need to zero out items in that list based on values in an index list, listB.

So, for example, given:

listA = [10, 12, 3, 8, 9, 17, 3, 7, 2, 8]
listB = [1, 4, 8, 9]

the output I want is

listC = [10, 0, 3, 8, 0, 17, 3, 7, 0, 0]

This question [1] seems similar, but asked for the elements to be removed, not changed. I'm not sure if a similar approach is needed, but if so I can't see how to apply it.

[1] how to remove elements from one list if other list contain the indexes of the elements to be removed

Community
  • 1
  • 1
MichaelMaggs
  • 260
  • 1
  • 11

2 Answers2

4

You can use a list comprehension, enumerate, and a conditional expression:

>>> listA = [10, 12, 3, 8, 9, 17, 3, 7, 2, 8]
>>> listB = [1, 4, 8, 9]
>>>
>>> list(enumerate(listA))  # Just to demonstrate
[(0, 10), (1, 12), (2, 3), (3, 8), (4, 9), (5, 17), (6, 3), (7, 7), (8, 2), (9, 8)]
>>>
>>> listC = [0 if x in listB else y for x,y in enumerate(listA)]
>>> listC
[10, 0, 3, 8, 0, 17, 3, 7, 0, 0]
>>>
Community
  • 1
  • 1
  • Many thanks for the help. I'd be happy to accept this, as it works fine, but as I can only accept one reply I've decided to go for mhlester's which also included an alternative approach. – MichaelMaggs May 01 '14 at 08:58
3

As a list comprehension:

listC = [value if index not in listB else 0 for index, value in enumerate(listA)]

Which for large lists can be improved substantially by using a set for listB:

setB = set(listB)
listC = [value if index not in setB else 0 for index, value in enumerate(listA)]

Or copy the list and modify it, which is both faster and more readable:

listC = listA[:]
for index in listB:
    listC[index] = 0
mhlester
  • 22,781
  • 10
  • 52
  • 75
  • I wouldn't be so sure the former is faster: it does a slow list membership test for each index and an `enumerate` call. The second approach does one fast copy and then only zeroes as many elements as it needs to. – DSM Apr 30 '14 at 19:01
  • Heh, you're probably right. It's funny how easy it is to fall into the trap of always assuming a list comprehension is better – mhlester Apr 30 '14 at 19:02
  • Thanks for the reply. In my actual data, listA has one million items and listB has 1089. Using %%timeit gives 27.5 seconds for the first approach and only 13ms for the second. Can that be right? It seems a huge difference. – MichaelMaggs May 01 '14 at 08:48
  • @MichaelMaggs: testing membership in lists is slow, because the elements have to be scanned one by one (and if an element isn't found, you don't know that until you've scanned the entire list). It should be much faster if you did `setB = set(listB)` and then used `not in setB` instead of `not in listB`. – DSM May 01 '14 at 12:08
  • Yes, that makes a big difference. It brings the time for the list comprehension from 27.5 seconds down to 191 ms. Still, for handling large lists the simple loop still wins on speed. – MichaelMaggs May 01 '14 at 13:12