Special indexing of two 1D numpy / torch arrays to produce another array
Looking for an efficient (vectorized) way in Numpy or PyTorch to achieve the array z when given input arrays x and y:
1D array x contains a list of increasing ID's, each of which repeats 1 or more times (not necessarily repeating for the same number of times for each ID). For example, [0 0 0 1 1 2 2 2 2]
1D array y of 0's and 1's. There is at least one element equal to &quot;1&quot; for each unique ID in x. For example, [1 1 0 1 1 0 0 1 0].
1D output array z which is equal to y, but keeping only the first occurrence of &quot;1&quot; in y per ID in x. The remaining elements of y for that ID should be set to &quot;0&quot;. So in the example, the result would be [1 0 0 1 0 0 0 1 0]
x: [0 0 0 1 1 2 2 2 2]
y: [1 1 0 1 1 0 0 1 0]
z: [1 0 0 1 0 0 0 1 0]

I feel like there's a quick way to do this in Numpy or PyTorch, but I couldn't figure it out.
Edit: Here's the &quot;slow&quot; version using a while-loop
x = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2])
y = np.array([1, 1, 0, 1, 1, 0, 0, 1, 0])
z = y.copy()
n = z.shape[0]
i = 0
while i &lt; n:
    if y[i] == 1:
        current_id = x[i]
        i += 1
        while i &lt; n and x[i] == current_id:
            z[i] = 0
            i += 1
    else:
        i += 1


You could use np.unique:
unq, ind = np.unique(np.stack((x, y)), axis=1, return_index=True)

ind now contains the first occurrence of each unique combination of elements. You just need to remove the ones where y is zero:
keep = unq[1, :] != 0
ind = ind[keep]

Now you can make z directly:
z = np.zeros_like(y)
z[ind] = 1

