python - form new list from n elements to the right of a reoccurring word

Given a list of strings:

haystack = ['hay','hay','hay','needle','x','y','z','hay','hay','hay','hay','needle','a','b','c']

Question

How would I form a new list of strings that contain, say, only the three adjacent elements (to the right) of every 'needle' occurrence within haystack?

Solution 1:

Find all the indices of "needle" and take 3 values right the indices.

# Get all indices of "needle"
idx = [idx for idx, val in enumerate(haystack) if val=="needle"]
#idx -> [3, 11]

# Take 3 values right of each index in `idx`.
[val for i in idx for val in haystack[i: i+4]] 
# ['needle', 'x', 'y', 'z', 'needle', 'a', 'b', 'c']

# want it to be a list of list
[haystack[i: i+4] for i in idx] 
# [['needle', 'x', 'y', 'z'], ['needle', 'a', 'b', 'c']]

# Want to exclude the "needle"
[val for i in idx for val in haystack[i+1: i+4]]
# ['x', 'y', 'z', 'a', 'b', 'c']

Solution 2:

This is a kind of hacky solution, but it works with only one pass through the list.

it = iter(haystack)
output = [[next(it), next(it), next(it)] for s in it if s == 'needle']
# [['x', 'y', 'z'], ['a', 'b', 'c']]

This is essentially the short-form of the following:

it = iter(haystack)
output = []
while True:
    try:
        elem = next(it)
        if elem == 'needle':
            output.append([next(it), next(it), next(it)])
    except StopIteration:
        break

note that, in the short form, you'll get a StopIteration error if there are fewer than three elements following a 'needle'.

Solution 3:

A simple list comprehension with list slicing seems to work as well:

out = [haystack[i+1:i+4] for i, x in enumerate(haystack) if x == 'needle']

Output:

[['x', 'y', 'z'], ['a', 'b', 'c']]