Special indexing of two 1D numpy / torch arrays to produce another array

user_14438074 0
Viewed 118

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 "1" 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 "1" in y per ID in x. The remaining elements of y for that ID should be set to "0". 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 "slow" 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 < n:
    if y[i] == 1:
        current_id = x[i]
        i += 1
        while i < n and x[i] == current_id:
            z[i] = 0
            i += 1
    else:
        i += 1
torch pytorch numpy arrays
2 Answers

If you really want to do this in a vectorized way, you could do this:

x = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2, 2])         
y = torch.tensor([1, 1, 0, 1, 1, 0, 0, 1, 0])         
                                                      
unique_vals = x.unique_consecutive().unsqueeze(1)     
masked_x = x.masked_fill(y != 1, unique_vals[-1][0]+1)
indices = (unique_vals == masked_x).int().argmax(1)   
res = torch.zeros_like(y).index_fill(0, indices, 1)   
print(res)

But I'm not convinced it'll be faster than the naive version.

user_32005540

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
user_29887300
Related