Index of first occurrence on sorted 1D tensor

tdigt: <p>Hello everyone!<br>
I’m trying to implement the following function as efficiently as possible:<br>
Given I have a sorted 1D tensor of positive integers, I want to retrieve the index of the first occurrence of each value. In case any integer between 0 and N-1 is missing, I want to return a -1. Thus, the function would return a 1D tensor of N elements.<br>
Is there an efficient way of doing this better than just a for loop over all values?</p>
<p>E.g.</p>
<pre><code class="lang-auto">x = torch.tensor([2, 3, 3, 4, 6, 6, 6, 8, 8], dtype=torch.long)
y = arg_first_ocurrence(x, N=10)
</code></pre>
<p>would return</p>
<pre><code class="lang-auto">&gt;&gt;&gt; y
tensor([-1, -1, 0, 1, 3, -1,  4, -1, 7, -1])
</code></pre>
<p>I’m open to solutions that would involve coding a C++/CUDA extension.</p>
<p>Thank you for your help!</p>

KFrank: <p>Hi tdigt!</p>
<aside class="quote no-group" data-username="tdigt" data-post="1" data-topic="113460" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tdigt/48/35285_2.png" class="avatar"> tdigt:</div>
<blockquote>
<p>E.g.</p>
<pre><code class="lang-auto">x = torch.tensor([2, 3, 3, 4, 6, 6, 6, 8, 8], dtype=torch.long)
y = arg_first_ocurrence(x, N=10)
</code></pre>
<p>would return</p>
<pre><code class="lang-auto">&gt;&gt;&gt; y
tensor([-1, -1, 0, 1, 3, -1,  4, -1, 7, -1])
</code></pre>
</blockquote>
</aside>
<p>I don’t know whether this meets your efficiency requirements, but<br>
here is a scheme that uses no loops, only pytorch tensor operations.</p>
<p>Note that <code>msk = x == lbl.unsqueeze (1)</code> uses broadcasting,<br>
and therefore generates a two-dimensional tensor from your<br>
one-dimensional input, possibly introducing inefficiency.</p>
<p>Also for convenience, I’ve eliminated the <code>N</code> argument, instead just<br>
using the length of the input tensor.</p>
<p>Here is the code in a small test script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch

torch.__version__

def arg_first_occurrence (x):
    n = x.numel()
    lbl = torch.arange (n)
    msk = x == lbl.unsqueeze (1)
    mna = (lbl + 1) * msk
    mnb = torch.where (mna != 0, mna, (n + 1) * torch.ones (n, 1).long())
    foa = mnb.min (dim = 1)[0] - 1
    fob = torch.where (foa != n, foa, -torch.ones (n).long())
    return fob

x = torch.tensor([2, 3, 3, 4, 6, 6, 6, 8, 8], dtype=torch.long)

y = arg_first_occurrence (x)

ychk = torch.tensor([-1, -1, 0, 1, 3, -1,  4, -1, 7], dtype = torch.long)

print ('torch.equal (y, tchk) =', torch.equal (y, ychk))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt;
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt;
&gt;&gt;&gt; def arg_first_occurrence (x):
...     n = x.numel()
...     lbl = torch.arange (n)
...     msk = x == lbl.unsqueeze (1)
...     mna = (lbl + 1) * msk
...     mnb = torch.where (mna != 0, mna, (n + 1) * torch.ones (n, 1).long())
...     foa = mnb.min (dim = 1)[0] - 1
...     fob = torch.where (foa != n, foa, -torch.ones (n).long())
...     return fob
...
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.tensor([2, 3, 3, 4, 6, 6, 6, 8, 8], dtype=torch.long)
&gt;&gt;&gt;
&gt;&gt;&gt; y = arg_first_occurrence (x)
&gt;&gt;&gt;
&gt;&gt;&gt; ychk = torch.tensor([-1, -1, 0, 1, 3, -1,  4, -1, 7], dtype = torch.long)
&gt;&gt;&gt;
&gt;&gt;&gt; print ('torch.equal (y, tchk) =', torch.equal (y, ychk))
torch.equal (y, tchk) = True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>