Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Selecting rows before and after rows of interest in Pandas

Let's say I have a time series dataframe with a categorical variable and a value:

In [4]: df = pd.DataFrame(data={'category': np.random.choice(['A', 'B', 'C', 'D'], 11), 'value': np.random.rand(11)}, index=pd.date_range('2015-04-20','2015-04-30'))

In [5]: df
Out[5]:
           category     value
2015-04-20        D  0.220804
2015-04-21        A  0.992445
2015-04-22        A  0.743648
2015-04-23        B  0.337535
2015-04-24        B  0.747340
2015-04-25        B  0.839823
2015-04-26        D  0.292628
2015-04-27        D  0.906340
2015-04-28        B  0.244044
2015-04-29        A  0.070764
2015-04-30        D  0.132221

If I'm interested in rows with category A, filtering to isolate them is trivial. But what if I'm interested in the n rows before category A's as well? If n=2, I'd like to see something like:

In [5]: df[some boolean indexing]
Out[5]:
           category     value
2015-04-20        D  0.220804
2015-04-21        A  0.992445
2015-04-22        A  0.743648
2015-04-27        D  0.906340
2015-04-28        B  0.244044
2015-04-29        A  0.070764

Similarly, what if I'm interested in the n rows around category A's? Again if n=2, I'd like to see this:

In [5]: df[some other boolean indexing]
Out[5]:
           category     value
2015-04-20        D  0.220804
2015-04-21        A  0.992445
2015-04-22        A  0.743648
2015-04-23        B  0.337535
2015-04-24        B  0.747340
2015-04-27        D  0.906340
2015-04-28        B  0.244044
2015-04-29        A  0.070764
2015-04-30        D  0.132221

Thanks!

like image 270
lenderson Avatar asked Feb 09 '17 21:02

lenderson


1 Answers

n rows around category A's:

In [223]: idx = df.index.get_indexer_for(df[df.category=='A'].index)

In [224]: n = 1

In [225]: df.iloc[np.unique(np.concatenate([np.arange(max(i-n,0), min(i+n+1, len(df)))
                                            for i in idx]))]
Out[225]:
           category     value
2015-04-20        D  0.220804
2015-04-21        A  0.992445
2015-04-22        A  0.743648
2015-04-23        B  0.337535
2015-04-28        B  0.244044
2015-04-29        A  0.070764
2015-04-30        D  0.132221

In [226]: n = 2

In [227]: df.iloc[np.unique(np.concatenate([np.arange(max(i-n,0), min(i+n+1, len(df)))
                                            for i in idx]))]
Out[227]:
           category     value
2015-04-20        D  0.220804
2015-04-21        A  0.992445
2015-04-22        A  0.743648
2015-04-23        B  0.337535
2015-04-24        B  0.747340
2015-04-27        D  0.906340
2015-04-28        B  0.244044
2015-04-29        A  0.070764
2015-04-30        D  0.132221
like image 185
MaxU - stop WAR against UA Avatar answered Oct 25 '22 16:10

MaxU - stop WAR against UA