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