X is a tensor of shape (N, D, P) such that for each n and d, X[n, d, :] is either all 0’s, or it contains a single block of consecutive 1’s and the rest are 0’s.

Y is a tensor of shape (N, P) such that each row is sorted increasing.

I want to create a tensor Z of shape (N, D, 2) such that

a. If i is the index of the first 1 in X[n, d, :], then Z[n, d, 0] = Y[n, i]

b. If j is the index of the last 1 in X[n, d, :], then Z[n, d, 1] = Y[n, j]

c. If X[n, d, :] is all 0’s, then I don’t care what Z[n, d, :] is.

Is there an efficient way to create Z?