python - 如何提高数据框中的搜索索引

Question

给定一个带有时间戳索引的 pandas 数据框，已排序。我有一个标签，我需要找到最接近该标签的索引。另外，我需要找到一个更小的时间戳，所以应该在次要时间戳中计算搜索。这是我的代码:

import pandas as pd
import datetime

data = [i for i in range(100)]
dates = pd.date_range(start="01-01-2018", freq="min", periods=100)
dataframe = pd.DataFrame(data, dates)

label = "01-01-2018 00:10:01"
method = "pad"
tol = datetime.timedelta(seconds=60)
idx = dataframe.index.get_loc(key=label, method="pad", tolerance=tol)

print("Closest idx:"+str(idx))
print("Closest date:"+str(dataframe.index[idx]))

搜索速度太慢。有什么办法可以改善吗？

Answer 1

为了提高性能，我建议对您搜索的内容进行转换。除了使用 get_loc，您还可以将 DateTimeIndex 转换为 Unix 时间，并使用 np.searchsorted在底层 numpy 数组上(顾名思义，这需要一个排序索引)。

---

get_loc:

(您当前的方法)

label = "01-01-2018 00:10:01"
tol = datetime.timedelta(seconds=60)
idx = dataframe.index.get_loc(key=label, method="pad", tolerance=tol)
print(dataframe.iloc[idx])

0    10
Name: 2018-01-01 00:10:00, dtype: int64

时间安排:

%timeit dataframe.index.get_loc(key=label, method="pad", tolerance=tol)
2.03 ms ± 81.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.searchsorted:

arr = df.index.astype(int)//10**9
l = pd.to_datetime(label).timestamp()
idx = np.max(np.searchsorted(arr, l, side='left')-1, 0)
print(dataframe.iloc[idx])

0    10
Name: 2018-01-01 00:10:00, dtype: int64

时间安排:

%timeit np.max(np.searchsorted(arr, l, side='left')-1, 0)
56.6 µs ± 979 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

(我没有包括设置成本，因为初始数组创建应该是你做一次，然后用于每个查询，但即使我确实包括了设置成本，这种方法也更快):

%%timeit
arr = df.index.astype(int)//10**9
l = pd.to_datetime(label).timestamp()
np.max(np.searchsorted(arr, l, side='left')-1, 0)

394 µs ± 3.84 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

上述方法不强制执行 60s 的容差，尽管这很容易检查:

>>> np.abs(arr[idx]-l)<60
True