这个question 提到了 numpy.r_ 但我不知道如何让它与切片列表一起工作,所以我使用了 hstack 和 arange
import numpy as np
import pandas as pd
def loop_version(df, desired):
# let's loop through the desired ranges and compile our selection
x = pd.DataFrame()
for (start, stop) in desired:
selection = df[(df.index >= pd.Timestamp(start)) &
(df.index <= pd.Timestamp(stop))]
x = x.append(selection)
# and let's have a look at what we found ...
return x
def vectorized_version(df, desired):
# first flatten the list
times = np.array(desired).flatten()
# use searchsorted to find the indices of the
# desired times in df's index
ndxlist = df.index.searchsorted(np.array(times))
# use np.arange to convert pairs of values in ndxlist to a
# range of indices, similar to np.r_
ndxlist = np.hstack([np.arange(i1, i2) for i1, i2 in
zip(ndxlist[::2], ndxlist[1::2])])
return df.iloc[ndxlist]
In [2]: # let's create some fake data
In [3]: date_range = pd.date_range('2005-01-01', '2008-12-31', freq='9min')
In [4]: l = len(date_range)
In [5]: df = pd.DataFrame({'normal': np.random.randn(l), 'uniform':np.random.rand(l),
...: 'datetime':date_range, 'integer':range(l)}, index=date_range)
In [6]: # let's identify the periods we want
...: desired = [('2005-10-27 14:30','2005-10-27 15:15'),
...: ('2006-04-14 14:40','2006-04-14 15:20'),
...: ('2008-01-25 14:30','2008-01-25 15:30')]
In [7]: loop_version(df, desired).equals(vectorized_version(df, desired))
Out[7]: True
In [8]: % timeit loop_version(df, desired)
5.53 ms ± 225 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [9]: % timeit vectorized_version(df, desired)
308 µs ± 1.26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)