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

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?

CodePudding user response：

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

CodePudding user response：

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