Numpy array: iterate through column and change value based on the current value and the next value-CodePudding

I have an array like this: This is an extension of a recent question that I asked elsewhere here. I have a numpy array like this:

data = np.array([
     [1,2,3],
     [1,2,3],
     [1,2,101],
     [4,5,111],   
     [4,5,6], 
     [4,5,6], 
     [4,5,101], 
     [4,5,112], 
     [4,5,6], 
     ])

In the third column, I want the value to be replaced with 10001 if the next one along is 101 AND if the current one is 6. which would result in an array like this:

data = np.array([
     [1,2,3],
     [1,2,3],
     [1,2,101],
     [4,5,111],   
     [4,5,6], 
     [4,5,10001], 
     [4,5,101], 
     [4,5,112], 
     [4,5,6], 
     ])

Any help on this would be greatly appreciated! Thanks!

CodePudding user response：

One way using numpy.roll:

s = data[:, 2]
data[np.logical_and(s == 6, np.roll(s, -1) == 101), 2] = 10001

Output:

array([[    1,     2,     3],
       [    1,     2,     3],
       [    1,     2,   101],
       [    4,     5,   111],
       [    4,     5,     6],
       [    4,     5, 10001],
       [    4,     5,   101],
       [    4,     5,   112],
       [    4,     5,     6]])

CodePudding user response：

i also came up with a working solution in the meantime:

data[
    
    (data[:,2] == 6)
    
    & 
      
      (np.roll(data[:,2]==101,-1,0)),2

] = 10001