Most computationally efficient way to batch alter values in each array of a 2d array, based on conditions for particular values by indices

Question

Say that I have a batch of arrays, and I would like to alter them based on conditions of particular values located by indices.

For example, say that I would like to increase and decrease particular values if the difference between those values are less than two.

For a single 1D array it can be done like this

import numpy as np

single2 = np.array([8, 8, 9, 10])

if abs(single2[1]-single2[2])<2:
    single2[1] = single2[1] - 1
    single2[2] = single2[2] + 1
single2

array([ 8,  7, 10, 10])

But I do not know how to do it for batch of arrays. This is my initial attempt

import numpy as np

single1 = np.array([6, 0, 3, 7])
single2 = np.array([8, 8, 9, 10])
single3 = np.array([2, 15, 15, 20])

batch = np.array([
    np.copy(single1),
    np.copy(single2),
    np.copy(single3),
])

if abs(batch[:,1]-batch[:,2])<2:
    batch[:,1] = batch[:,1] - 1
    batch[:,2] = batch[:,2] + 1

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Looking at np.any and np.all, they are used to create an array of booleans values, and I am not sure how they could be used in the code snippet above.

My second attempt uses np.where, using the method described here for comparing particular values of a batch of arrays by creating new versions of the arrays with values added to the front/back of the arrays.

https://stackoverflow.com/a/71297663/3259896

In the case of the example, I am comparing values that are right next to each other, so I created copies that shift the arrays forwards and backwards by 1. I also use only the particular slice of the array that I am comparing, since the other numbers would also be used in the comparison in np.where.

batch_ap = np.concatenate(
    (batch[:, 1:2+1], np.repeat(-999, 3).reshape(3,1)), 
    axis=1
)

batch_pr = np.concatenate(
    (np.repeat(-999, 3).reshape(3,1), batch[:, 1:2+1]), 
    axis=1
)

Finally, I do the comparisons, and adjust the values

batch[:, 1:2+1] = np.where(
    abs(batch_ap[:,1:]-batch_ap[:,:-1])<2,
    batch[:, 1:2+1]-1,
    batch[:, 1:2+1]
)

batch[:, 1:2+1] = np.where(
    abs(batch_pr[:,1:]-batch_pr[:,:-1])<2,
    batch[:, 1:2+1]+1,
    batch[:, 1:2+1]
)
print(batch)

[[ 6  0  3  7]
 [ 8  7 10 10]
 [ 2 14 16 20]]

Though I am not sure if this is the most computationally efficient nor programmatically elegant method for this task. Seems like a lot of operations and code for the task, but I do not have a strong enough mastery of numpy to be certain about this.

Answer 1

This works

mask = abs(batch[:,1]-batch[:,2])<2
batch[mask,1] -= 1
batch[mask,2] += 1