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

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. which would result in an array like this:

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

I tried this which I was certain would work but it doesn't...

dith = np.nditer(data[:, 2], op_flags=['readwrite'])
for i in dith:
    if i 1 == 3:
        i[...] = 10001

If anyone could help with this then that would be great.

CodePudding user response：

Try this -

data[np.roll(data==101,-1,0)] = 10001

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

Only assumption here is that your first row doesn't contain a 101

In case there is a potential scenario that 101 may occur in the first row of the matrix, then try this approach below.

idx = np.vstack([np.roll(data==101,-1,0)[:-1], np.array([False, False, False])])
data[idx] = 10001

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

CodePudding user response：

indices_of_101 = np.where(data[:, 2] == 101)[0]
if indices_of_101[0] = 0:  # taking into accound boundary problem
    indices_of_101 = indices_of_101[1:]
data[:, indices_of_101-1] = 10001