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?
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])