kalpesh-interview-question.py

kalpesh-interview-question.py

"""
How to find intersection of big tensor and all small tensors?
"""

import torch

bigsize = 30
smallsize = 8
smallcount = 12
vocabsize = 100
seed = 10

torch.manual_seed(10)

big = torch.LongTensor(1, bigsize).random_(0, vocabsize)
small = torch.LongTensor(smallcount, smallsize).random_(0, vocabsize)

print('SETUP')
print('big')
print(big)

print('small')
print(small)
print()

bigshape = big.shape
smallshape = small.shape
small = small.view(-1)
big = big.view(-1)

bigindex = big.argsort()
smallindex = small.argsort()
biginvert = bigindex.argsort()
smallinvert = smallindex.argsort()
bigvec = big[bigindex]
smallvec = small[smallindex]

bigmask = torch.zeros(*big.shape)
smallmask = torch.zeros(*small.shape)

j = 0
for i in range(small.shape[0]):
    if smallvec[i].item() < bigvec[j].item():
        continue
    while smallvec[i] > bigvec[j]:
        j += 1
        if j >= big.shape[0]:
            break
        if bigvec[j].item() == bigvec[j-1].item():
            bigmask[j] = bigmask[j-1]
    if j >= big.shape[0]:
        break

    if smallvec[i].item() == bigvec[j].item():
        smallmask[i] = 1
        bigmask[j] = 1

# Cleanup any remaining elements of the bigger tensor's mask.
while j < big.shape[0]:
    j += 1
    if j >= big.shape[0]:
        break
    if bigvec[j].item() == bigvec[j-1].item():
        bigvec[j] = bigvec[j-1]

bigmask = bigmask[biginvert]
smallmask = smallmask[smallinvert]

big = big.view(*bigshape)
small = small.view(*smallshape)
smallmask = smallmask.view(*smallshape)
bigmask = bigmask.view(*bigshape)

print('RESULT')
print('big')
print(big)

print('small')
print(small)

print('bigmask')
print(bigmask)

print('smallmask')
print(smallmask)

I do all small vectors at once, but it might be faster if you do them all separately (less things to sort).