Find the indices at which any element of one list

Question

Take lists haystack and needles

haystack = ['a', 'b', 'c', 'V', 'd', 'e', 'X', 'f', 'V', 'g', 'h']
needles = ['V', 'W', 'X', 'Y', 'Z']

I need to generate a list of the indices at which any element of needles occurs in haystack. In this case those indices are 3, 6, and 8 thus

result = [3, 6, 8]

This question I found is very similar and was rather elegantly solved with

result = [haystack.index(i) for i in needles]

Unfortunately, this solution gives ValueError: 'W' is not in list in my case. This is because the difference here is that an element of needles may occur in haystack a number of times or not at all.

In other words, haystack may contain no needles or it may contain many.

Answer 1

In addition to failing if your needle is not in the haystack, the index method will return only the first position of the element you're looking for, even if that element appears more than once (as in the 'V' in your example). You could just do this:

result = [idx for idx, val in enumerate(haystack) if val in needles]

The enumerate function produces a generator that yields tuples of values - the first being the index and the second being a value:

>>> print(list(enumerate(['a', 'b', 'c'])))

Just check if each value is in your needles list and add the index if it is.

Answer 2

haystack = ['a', 'b', 'c', 'V', 'd', 'e', 'X', 'f', 'V', 'g', 'h']
needles = ['V', 'W', 'X', 'Y', 'Z']
st = set(needles)
print([i for i, e in enumerate(haystack) if e in st])
[3, 6, 8]

Even if you used [haystack.index(i) for i in needles if i in haystack] it would not work as you have repeated elements.

Making st = set(needles) means we have a linear solution as set lookups are 0(1) which for large input would be significantly more efficient.

Answer 3

needles_set = set(needles)
print [i for i, val in enumerate(haystack) if val in needles_set]

Answer 4

Definitely not the most efficient way but you could do something like this:

result = []
i=0
while (i < len(haystack)):
    if (needles.count(haystack[i]) > 0):
        result.append(i)
    i+=1

Which will make result = [3, 6, 8]

Answer 5

You could try something like the following.

[Haystack.index(x) for x in needles if x in Haystack]

If x is not in haystack then haystack.index(x) will not be called and no error should be thrown.