Efficient way for calculating selected differences in array
<p>I have two arrays as an output from a simulation script where one contains IDs and one times, i.e. something like:</p>

<pre><code>ids = np.array([2, 0, 1, 0, 1, 1, 2])
times = np.array([.1, .3, .3, .5, .6, 1.2, 1.3])
</code></pre>

<p>These arrays are always of the same size. Now I need to calculate the differences of <code>times</code>, but only for those times with the same <code>ids</code>. Of course, I can simply loop over the different <code>ids</code> an do</p>

<pre><code>for id in np.unique(ids):
    diffs = np.diff(times[ids==id])
    print diffs
    # do stuff with diffs
</code></pre>

<p>However, this is quite inefficient and the two arrays can be very large. Does anyone have a good idea on how to do that more efficiently?</p>

<p>Say you <code>np.argsort</code> by <code>ids</code>:</p>

<pre><code>inds = np.argsort(ids, kind='mergesort')
&gt;&gt;&gt; array([1, 3, 2, 4, 5, 0, 6])
</code></pre>

<p>Now sort <code>times</code> by this, <code>np.diff</code>, and prepend a <code>nan</code>:</p>

<pre><code>diffs = np.concatenate(([np.nan], np.diff(times[inds])))
&gt;&gt;&gt; diffs 
array([ nan,  0.2, -0.2,  0.3,  0.6, -1.1,  1.2])
</code></pre>

<p>These differences are correct except for the boundaries. Let's calculate those</p>

<pre><code>boundaries = np.concatenate(([False], ids[inds][1: ] == ids[inds][: -1]))
&gt;&gt;&gt; boundaries
array([False,  True, False,  True,  True, False,  True], dtype=bool)
</code></pre>

<p>Now we can just do</p>

<pre><code>diffs[~boundaries] = np.nan
</code></pre>

<hr>

<p>Let's see what we got:</p>

<pre><code>&gt;&gt;&gt; ids[inds]
array([0, 0, 1, 1, 1, 2, 2])

&gt;&gt;&gt; times[inds]
array([ 0.3,  0.5,  0.3,  0.6,  1.2,  0.1,  1.3])

&gt;&gt;&gt; diffs
array([ nan,  0.2,  nan,  0.3,  0.6,  nan,  1.2])
</code></pre>

python arrays numpy