Skip to content
Advertisement

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

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!

Advertisement

Answer

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]])
User contributions licensed under: CC BY-SA
3 People found this is helpful
Advertisement