Get next value from a row that satisfies a condition in pandas
<p>I have a DataFrame that looks something like this:</p>

<pre><code>       | event_type | object_id
------ | ------     | ------
0      | A          | 1
1      | D          | 1
2      | A          | 1
3      | D          | 1
4      | A          | 2
5      | A          | 2
6      | D          | 2
7      | A          | 3
8      | D          | 3
9      | A          | 3
</code></pre>

<p>What I want to do is get the index of the next row where the <code>event_type</code> is A and the <code>object_id</code> is still the same, so as an additional column this would look like this:</p>

<pre><code>       | event_type | object_id | next_A
------ | ------     | ------    | ------
0      | A          | 1         | 2
1      | D          | 1         | 2
2      | A          | 1         | NaN
3      | D          | 1         | NaN
4      | A          | 2         | 5
5      | A          | 2         | NaN
6      | D          | 2         | NaN
7      | A          | 3         | 9
8      | D          | 3         | 9
9      | A          | 3         | NaN
</code></pre>

<p>and so on.</p>

<p>I want to avoid using <code>.apply()</code> because my DataFrame is quite large, is there a vectorized way to do this?</p>

<p>EDIT: for multiple A/D pairs for the same <code>object_id</code>, I'd like it to always use the next index of A, like this:</p>

<pre><code>       | event_type | object_id | next_A
------ | ------     | ------    | ------
0      | A          | 1         | 2
1      | D          | 1         | 2
2      | A          | 1         | 4
3      | D          | 1         | 4
4      | A          | 1         | NaN
</code></pre>

<p>You can do it with groupby like:</p>

<pre><code>def populate_next_a(object_df):
    object_df['a_index'] = pd.Series(object_df.index, index=object_df.index)[object_df.event_type == 'A']
    object_df['a_index'].fillna(method='bfill', inplace=True)
    object_df['next_A'] = object_df['a_index'].where(object_df.event_type != 'A', object_df['a_index'].shift(-1))
    object_df.drop('a_index', axis=1)
    return object_df
result = df.groupby(['object_id']).apply(populate_next_a)
print(result)
  event_type  object_id  next_A
0          A          1     2.0
1          D          1     2.0
2          A          1     NaN
3          D          1     NaN
4          A          2     5.0
5          A          2     NaN
6          D          2     NaN
7          A          3     9.0
8          D          3     9.0
9          A          3     NaN
</code></pre>

<p><a href="http://pandas.pydata.org/pandas-docs/stable/groupby.html" rel="nofollow">GroupBy.apply</a> will not have as much overhead as a simple apply.</p>

<p>Note you cannot (yet) store integer with NaN: <a href="http://pandas.pydata.org/pandas-docs/stable/gotchas.html#support-for-integer-na" rel="nofollow">http://pandas.pydata.org/pandas-docs/stable/gotchas.html#support-for-integer-na</a> so they end up as float values</p>

python python-3.x pandas