Find a list of items and return x number of surrounding items in python
I want to find a list for the occurrence of a value (x) and return that value and a number, say 2, of values above and below x in the index. The x value can appear multiple times in the list.
Input
in = ['a','b','c','d','x','e','f','g','h','i','x','j','k','l']
Output
out = ['c','d','x','e','f','h','i','x','j','k']
Thanks for any help or suggestions
+3
source to share
3 answers
In [8]: lis = ['a','b','c','d','x','e','f','g','h','i','x','j','k','l']
#create a new list containing all the index positions of 'x'
In [9]: ind=[i for i,x in enumerate(lis) if x=='x']
In [10]: out=[]
# loop over ind list, and for every index i:
# here lis[i-2:i] are the elements left to the 'x' and similarly lis[i:i+3]
# are the ones to its right.
# which is simply lis[i-2:i+3] as suggested by @volatility
In [11]: for i in ind:
out.extend(lis[i-2:i+3])
....:
In [12]: out
Out[12]: ['c', 'd', 'x', 'e', 'f', 'h', 'i', 'x', 'j', 'k']
One-liner using itertools.chain()
:
In [19]: from itertools import *
In [20]: list(chain(*[lis[i-2:i+3] for i in ind]))
Out[20]: ['c', 'd', 'x', 'e', 'f', 'h', 'i', 'x', 'j', 'k']
+3
source to share
Alternative solution using difflib.SequenceMatcher
>>> from itertools import chain
>>> from difflib import SequenceMatcher
>>> in_data = ['a','b','c','d','x','e','f','g','h','i','x','j','k','l']
>>> sm = SequenceMatcher(None, in_data, 'x'*len(in_data)).get_matching_blocks()
>>> list(chain(*(in_data[m.a -2 : m.a + 3] for m in sm[:-1])))
['c', 'd', 'x', 'e', 'f', 'h', 'i', 'x', 'j', 'k']
+1
source to share
l = ['a','b','c','d','x','e','f','g','h','i','x','j','k','l']
def search_w(mylist,item,before=1,after=1):
newl=[]
l = mylist[:]
while item in l:
i = l.index(item)
newl+= l[i-before:i+after+1]
l = l[i+after:]
return newl
>>> print search_w(l,'x',2,2)
['c', 'd', 'x', 'e', 'f', 'h', 'i', 'x', 'j', 'k']
+1
source to share