How to split a numpy array of integers, into chunks that have successive values with a difference below a threshold

Question

I have the following numpy array with positive integers, in ascending order:

import numpy as np

arr = np.array([222, 225, 227, 228, 230, 232, 241, 243, 244, 245, 252, 253, 258])

I want to split it, into parts, where at each part, each number has maximum difference of 2 from the next one.

So the following array should be split as:

[[222], [225,227,228,230,232], [241, 243, 244, 245], [252, 253], [258]]

Any ideas how I can I achieve that ?

Answer 1

You can compute the diff, get the indices of differences above threshold with flatnonzero, and split with array_split:

threshold = 2

out = np.array_split(arr, np.flatnonzero(np.diff(arr)>threshold)+1)

output:

[array([222]),
 array([225, 227, 228, 230, 232]),
 array([241, 243, 244, 245]),
 array([252, 253]),
 array([258])]