group list elements by their difference with respect to each other

Question

Given a list

A = [1, 6, 13, 15, 17, 18, 19, 21, 29, 36, 53, 58, 59, 61, 63, 78, 79, 81, 102, 114]

is there an easy way to group all those clusters where the difference between successive elements is smaller than 3?

That is, obtain something like:

[13, 15, 17, 19, 21], [58, 59, 61, 63], [78, 79, 81]

I was wondering whether there exists any built-in function, but I couldn't find anything similar. I was trying to figure it out using the groupby of the itertools, but I'm stuck. Thank you in advance.

Answer 1

You can use itertools.groupby:

import itertools
A = [1, 6, 13, 15, 17, 18, 19, 21, 29, 36, 53, 58, 59, 61, 63, 78, 79, 81, 102, 114]
new_a = [(A[i+1]-A[i], A[i]) for i in range(len(A)-1)]
a = [[a, [c for _, c in b]] for a, b in itertools.groupby(new_a, key=lambda x:x[0] < 3)]
final_groups = [a[i][-1]+[a[i+1][-1][0]] if a[i+1][-1][0] - a[i][-1][-1] < 3 else a[i][-1] for i in range(len(a)-1) if a[i][0]]

Output:

[[13, 15, 17, 18, 19, 21], [58, 59, 61, 63], [78, 79, 81]]

Answer 2

This is one approach using an iteration.

Ex:

A = [1, 6, 13, 15, 17, 18, 19, 21, 29, 36, 53, 58, 59, 61, 63, 78, 79, 81, 102, 114]
res = []
temp = []
l = len(A)-1

for i,v in enumerate(A):
    if i+1 > l:
        break

    if abs(v - A[i+1]) < 3:
        temp.append(v)
    else:
        if temp:
            temp.append(v)
            res.append(temp)
            temp = []
print(res)

Output:

[[13, 15, 17, 18, 19, 21], [58, 59, 61, 63], [78, 79, 81]]