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

Question

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?

Answer 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']