Replace the item and the following items in the list if the condition is met

Question

I have a list that looks like this:

a=[1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1]

I am looking for a way that every time it finds 1, this element AND the next two elements are replaced with -999. The following code seems to work.

b = a
x = 2 #number of continuous elements to browse 
for n,i in enumerate(a):
    if i==1:
        for t in range(x):
            print n+t
            b[n+t]=-999

print b

However, when 1 is at the end of the list, I get:

IndexError: list destination index out of range

Is there a better way to do this?

Answer 1

This is a special case - for elements at the end of the list, the next element does not exist. You must check this special case in the inner loop before you access the list:

if n+t >= len(b):
    break

Answer 2

Single line:

b = [(i, -999)[1 in a[max(0,n-x):n+1]] for n, i in enumerate(a)]

It produces:

# x = 2
[-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, -999, -999, -999, 3, 2, -999, -999, -999]
# x = 1
[-999, -999, 3, 4, 5, 2, 3, -999, -999, 3, 4, 5, -999, -999, 5, 3, 2, -999, -999, -999]

(i, -999)[1 in a[max(0,n-x):n+1]] : -999, 1 x, i.

Answer 3

seen, 0, 1, 2:

def replace_element_and_next(seq, search, replace):
    seen = 0
    for x in seq:
        if x == search:
            yield replace
            seen = 1
        elif seen in [1, 2]:
            yield replace
            seen = (seen + 1) % 3
        else:
            yield x

a = [1, 2, 3, 4, 5, 2, 3, 1, 2, 3, 4, 5, 1, 5, 5, 3, 2, 1, 3, 1]
print(','.join('%4d' % x for x in a))
b = replace_element_and_next(a, 1, -999)
print(','.join('%4d' % x for x in b))

:

   1,   2,   3,   4,   5,   2,   3,   1,   2,   3,   4,   5,   1,   5,   5,   3,   2,   1,   3,   1
-999,-999,-999,   4,   5,   2,   3,-999,-999,-999,   4,   5,-999,-999,-999,   3,   2,-999,-999,-999

• seen 0
• seen (0, 1, 2) , 0, 1 2
• seen = (seen + 1) % 3 , seen 0, 1 2

• replace_element_and_next() -. , :

  b = list(replace_element_and_next(a, 1, -999))
  

Answer 4

- itertools.groupby:

import itertools
a=[1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1]
new_a = [(a, list(b)) for a, b in itertools.groupby(a, key=lambda x:x == 1)]
final_a = list(itertools.chain(*[[-999] if a else [-999, -999, *b[2:]] for a, b in new_a]))

:

[-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, -999, -999, -999, 3, 2, -999, -999, -999, -999]

: Python2:

import itertools
a=[1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1]
new_a = [(a, list(b)) for a, b in itertools.groupby(a, key=lambda x:x == 1)]
final_a = list(itertools.chain(*[[-999] if a else [-999, -999]+b[2:] for a, b in new_a]))

Answer 5

, , , python , . try , .

b = a
x = 2 #number of continuous elements to browse 
for n,i in enumerate(a):
    if i==1:
        for t in range(x):
            print n+t
            try: b[n+t]=-999
            except IndexError:
                # either pass or append right? 
                b.append(-999)
                # pass
print b

, , , .

Answer 6

, , , , , , ,

:

, "1" -999 , [-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, 1, -999, -999, 3, 2, -999, -999, -999, -999]

:

, - :

a=[1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1,1]


for i,j in enumerate(a):

    try:
        if j == 1:
            a[i] = -999
            a[i + 1] = -999
            a[i + 2] = -999

    except IndexError:
        pass

print(a)

:

[-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, -999, -999, -999, 3, 2, -999, -999, -999, -999]

Answer 7

, indices , 1 a, , , , , a[21], 20. , , a, IndexError: list assignment index out of range.

, a, .

:

a=[1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1]

indices = set()
for i, e in enumerate(a):
    if e == 1:
        indices.add(i)
        indices.add(i+1)
        indices.add(i+2)

pos = {x for x in indices if x < len(a)}

for i in pos:
    a[i] = -999

print(a)

:

[-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, -999, -999, -999, 3, 2, -999, -999, -999]

.. , O(N). , .

Answer 8

Using range()and len(), you can try the following:

a = [1,2,3,4,5,2,3,1,2,3,4,5,1,5,5,3,2,1,3,1]

for i in range(len(a)):
    if a[i] == 1:
        if len(a) - i >= 2:
            a[i] = -999
            a[i+1] = -999
            a[i+2] = -999

print(a)

[-999, -999, -999, 4, 5, 2, 3, -999, -999, -999, 4, 5, -999, -999, -999, 3, 2, -999, -999, -999]
>>>