I have two tensors in pytorch:

tensorA=[0, 1, 2, 6, 7, 9, 10]


tensorB=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

How can I use pytorch's tensor APIs(ops) to find the elements in tensorB that are missing from TensorA, then for each missing element, return the index of the nearest previous element in TensorA?

For example, for the above tensors, it should return:

[2, 2, 2, 4] 

because [3, 4, 5], and [8] are missing from tensorA. For [3, 4, 5], we need to return their previous neighbouring element's index in tensorA, which is 2. For element [8], its previous neighbouring element in tensorA is 7, and its index in tensorA is 4, so we return the total tensor as: [2, 2, 2, 4].

Can you provide pytorch's torch APIs to implement the above function? I need to fully utilize the vectorized APIs to accelerate my function, so can't simply use a loop. Thanks!


1 Answer 1


Seeing your tensor B is just a range of values, I assumed that tensor B always covers the full range of 0 to n.

In addition, you did not specify what had to happen if the first couple values were missing (as then there is no smaller value, so i left that out.

Here is the code:

begin_index = torch.tensor([0])
end_index = torch.tensor([11])

a = torch.tensor([0, 1, 2, 6, 7, 9, 10])

missing = a[1:] - (a[:-1] + 1)
vals = missing.nonzero().squeeze()
times = missing[missing.nonzero()].squeeze()
indices = vals.repeat_interleave(times)
end = torch.tensor([a.shape[0]-1]) * (end_index - a[-1])
final = torch.cat((indices, end))

begin_missing = (a[0] - begin_index).item()
print("Missing " + str(begin_missing) + " items at the start of the list")

Which outputs

Missing 0 items at the start of the list
tensor([2, 2, 2, 4, 6])

The 6 is added because i put the end_index on 11, in case the last indices are also not included. I'm not 100% sure whether this is to your satisfaction, but I think it will go a long way.

Again, I removed Tensor B and simply assumed it covered the range of values from begin_index to end_index.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.