18

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.

Community
  • 1
  • 1
chambln
  • 363
  • 1
  • 2
  • 9

5 Answers5

23
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.

Padraic Cunningham
  • 176,452
  • 29
  • 245
  • 321
4
needles_set = set(needles)
print [i for i, val in enumerate(haystack) if val in needles_set]
Saksham Varma
  • 2,122
  • 13
  • 15
1

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.

Daniel Robertson
  • 1,354
  • 13
  • 22
0

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.

paidhima
  • 2,312
  • 16
  • 13
0

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]

Jon Warren
  • 857
  • 6
  • 18