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