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

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

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

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

If I understood correctly then you want this...

for i in [i for i,ele in enumerate(haystack) if ele=="needle"]:
    Out.extend(haystack[i+1:i+4])
print(Out)

Output

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