Finding indices of runs of integers in an array

user_6997665 0
Viewed 64

I have a numpy array with integers as entries. I need to find the indices corresponding to runs of integers in the array. For example if my array is a = [1, 2, 3, 6, 7, 8], then the indices I want are [[0, 1, 2], [3, 4, 5]]. I have tried the following piece of code which uses a loop.

import numpy as np

idx = np.array([0, 1, 2, 3, 8, 9, 10, 11, 12, 2, 3, 4, 5, 6, 55, 56, 89, 90, 91, 92])
idx_cpy = idx

indices = []

while len(idx) > 0:
    sorted_idx = np.arange(idx[0], idx.max()+1)[:len(idx)]
    bool_equal = np.equal(sorted_idx, idx)
    true_idx = np.argwhere(bool_equal == True)[:, 0]
    try:
        indices.append(np.array(indices[-1]).max() + 1 + true_idx)
    except IndexError:
        indices.append(true_idx)
    idx = idx[true_idx.max()+1:]


"""
indices = 
[array([0, 1, 2, 3], dtype=int64), 
array([4, 5, 6, 7, 8], dtype=int64), 
array([ 9, 10, 11, 12, 13], dtype=int64), 
array([14, 15], dtype=int64), 
array([16, 17, 18, 19], dtype=int64)]
"""

Although this works as expected, the variable idx in my actual code has large length and this takes a long time to complete. Is there a vectorized way of doing this? Thanks.

python-3.x numpy arrays python
1 Answers

You can find the "turning points" with np.diff and check where it differs from 1. To include end points, we pass prepend and append to it such that the difference isn't 1 in those places and they get counted too. Then a list comprehension with np.arange gives the final result:

>>> turnings, = np.where(np.diff(a, prepend=a[0], append=a[-1]) != 1)
>>> turnings

array([ 0,  4,  9, 14, 16, 20], dtype=int64)

>>> result = [np.arange(pre, nex) for pre, nex in zip(turnings, turnings[1:])]
>>> result

[array([0, 1, 2, 3], dtype=int64),
 array([4, 5, 6, 7, 8], dtype=int64),
 array([ 9, 10, 11, 12, 13], dtype=int64),
 array([14, 15], dtype=int64),
 array([16, 17, 18, 19], dtype=int64)]
user_93321870
Related