pythonalgorithmnumpy

Fast way to expand split list into index list


Given an index split list T of length M + 1, where the first element is 0 and the last element is N, generate an array D of length N such that D[T[i]:T[i+1]] = i.

For example, given T = [0, 2, 5, 7], then return D = [0, 0, 1, 1, 1, 2, 2].

I'm trying to avoid a for loop, but the best I can do is:

def expand_split_list(split_list):
    return np.concatenate(
        [
            np.full(split_list[i + 1] - split_list[i], i)
            for i in range(len(split_list) - 1)
        ]
    )

Is there a built-in function for that purpose?


Solution

  • You could combine diff, arange, and repeat:

    n = np.diff(T)
    out = np.repeat(np.arange(len(n)), n)
    

    As a one-liner (python ≥3.8):

    out = np.repeat(np.arange(len(n:=np.diff(T))), n)
    

    Another option with assigning ones to an array of zeros, then cumsum:

    out = np.zeros(T[-1], dtype=int)
    out[T[1:-1]] = 1
    out = np.cumsum(out)
    

    Output:

    array([0, 0, 1, 1, 1, 2, 2])