pandas dataframe look ahead optimization
<p>What's the best way to do this with pandas dataframe? I want to loop through a dataframe, and find the nearest next index that has at least +/-2 value difference. For example: [100, 99, 102, 98, 103, 103] will create a new column with this [2, 2, 3, 0, N/A], 0 means not find.</p>

<p>My solution performance is n * log(n). Any brilliant people could please show me a better performance solution?</p>

<p>When all the elements are integers, it is possible to do so in linear time. The following solution is complicated, and of algorithmic interest (if at that) only. Since it uses loops and data structures, any real implementation would need to be in C/C++/Cython (otherwise, the constants would be so high, that a tremendously long sequence would be needed to start seeing an improvement, even though it's linear).</p>

<p>Since the solution is complicated, I'll first make some simplifying assumptions, then show how to get rid of them. The initial assumptions are:</p>

<ol>
<li><p>What's needed is to find the index of the next position that is 2 or greater.</p></li>
<li><p>All the integers are distinct.</p></li>
</ol>

<p>Given these assumptions, it's possible to use a variant of a well known interview question (it's so common, I think it's folklore). The idea is to keep a stack of the positions of the array where the next positions have not yet been found. While looping over the elements and the positions, loop invariants are  maintained: </p>

<ul>
<li><p>The indices in the stack are increasing.</p></li>
<li><p>The stack does not contain positions <em>i</em>, <em>j</em> such that <em>a[i] + 2 &lt;= a[j]</em> and <em>i &lt; j</em>. </p></li>
</ul>

<p>The invariants are initially trivially satisfied, and I'll show how they're maintained.</p>

<p>Say at iteration <em>j</em>, the stack's top is <em>i</em>: I'll mark this as <em>&lt;..., i></em> (the stack goes to the right). While <em>a[j] >= a[i] + 2</em>, we can pop the stack and set the next position of <em>i</em> to be <em>j</em>. While this happens, we can pop the stack until the condition fails. As some point, though, the stack can be <em>&lt;..., k, i></em>, with <em>a[i] + 2 > a[j]</em>. Some thought on the invariants is enough to see, that in this case, if there is an element in the stack that needs to be popped, it must be <em>k</em> (if it exists). That is the only item that needs to be checked - any other item before the last one cannot be one that needs to be popped. So, we just need to check <em>k</em>, and pop it too if necessary. At the end of the iteration, we just need to push <em>j</em> itself.</p>

<p>The following code does this:</p>

<pre><code>def plus2_detector(a, verbose=False):
    if verbose:
        print 'starting with', a
    remaining, out = [], [None] * len(a)
    for i, e in enumerate(a):
        if verbose:
            print 'it begin', i, e, remaining
        while remaining and e &gt;= a[remaining[-1]] + 2:
            if verbose:
                print 'setting', i, remaining[-1], a[remaining[-1]]
            out[remaining[-1]] = i 
            del remaining[-1]
        if len(remaining) &gt; 1 and e &gt;= a[remaining[-2]] + 2:
            if verbose:
                print 'back setting', i, remaining[-2], a[remaining[-2]]
            out[remaining[-2]] = i 
            del remaining[-2]
        remaining.append(i) 
        if verbose:
            print 'it end', i, e, remaining
    return out
</code></pre>

<p>You can run it, e.g., </p>

<pre><code>&gt;&gt;&gt; plus2_detector([1, 2, 3, 5, 4, -1, -2, 10, 9, 8, 7, 11], False)
[2, 3, 3, 7, 7, 7, 7, None, 11, 11, 11, None]
</code></pre>

<p>To get an intuitive feeling of what it does, you can run it on different (distinct integers!) with <code>verbose=True</code>, and see what it does.</p>

<hr>

<p>Now to get rid of the simplifications. </p>

<p>The first simplification riddance is easy: run two copies of this algorithm: one checking for <em>>= 2</em>, and one checking for <em>&lt;= -2</em>, and combine the results.</p>

<p>The second simplification riddance is more tricky. The problem is that if the top of the stack doesn't need to be popped, we might need to search many items back to see if anyone needs to be popped - it's not necessarily true that this potential item is right under the top. This can happen if the elements along the top of the stack are identical.</p>

<p>Dealing with this problem is tedious, but not that difficult conceptually. The stack now needs to contain lists of integers of consecutive undealt indices of equivalent elements. This means that when you push a new index, you need to check if it continues a run. If it does, append it to the list at the top; if it does not, create a new tuple. Now all consecutive equivalent undealt items are grouped together (similar to what <code>itertools.groupby</code> does). </p>

<p>There are technical complications (when popping the penultimate list, we might need to combine the top and the new penultimate tuple), but the idea is the same.</p>

<p>The complexity is linear using a standard argument from amortized analysis (each element is inserted and popped once, the non-popping operations are constants).</p>

<p>Here is the code for the general case of finding +2 or above indices, without the restriction that elements are unique:</p>

<pre><code>def general_plus2_detector(a, verbose=False):
    if verbose:
        print 'starting with', a
    remaining, out = [], [None] * len(a)
    for i, e in enumerate(a):
        if verbose:
            print 'it begin', i, '(', e, ')', remaining
        while remaining and e &gt;= a[remaining[-1][0]] + 2:
            for j in remaining[-1]:
                if verbose:
                    print 'setting', j, '(', a[j], ') to', i, '(', a[i], ')'
                out[j] = i
            del remaining[-1]
        if len(remaining) &gt; 1 and e &gt;= a[remaining[-2][0]] + 2:
            for j in remaining[-2]:
                if verbose:
                    print 'back setting', j, '(', a[j], ') to', i, '(', a[i], ')'
                out[j] = i
            del remaining[-2]
            if len(remaining) &gt; 1 and a[remaining[-2][0]] == a[remaining[-1][0]]:
                if verbose:
                    print 'joining', remaining[-2], remaining[-1]
                remaining[-1].extend(remaining[-2])
                del remaining[-2]
        if not remaining or a[remaining[-1][0]] != e:
            remaining.append([i]) 
        else:
            remaining[-1].append(i)
        if verbose:
            print 'it end', i, '(', e, ')', remaining
    return out
</code></pre>

<p>Running it shows:</p>

<pre><code>a = [-1, -2, 3, 2, 2, 3, 2, 2, 4, 5, 4, 5, 2, 3, 4, 5, 5, 4, 4, 7]
&gt;&gt;&gt; general_plus2_detector(a, False)
[2, 2, 9, 8, 8, 9, 8, 8, 19, 19, 19, 19, 14, 15, 19, 19, 19, 19, 19, None]
</code></pre>

python algorithm