时间戳

首页 > 解决方案 > 优化滚动窗口熊猫数据框

python - 优化滚动窗口熊猫数据框

问题描述

我有一个这样的熊猫系列：

dist
0.02422
0.03267
0.04208
0.05229
0.06291
...

它有近 200K 行。是否有更有效的方法来执行以下操作：

df["dist"].rolling(3000).apply(lambda x: x.iloc[-1] - x[x != 0].iloc[0] if x[x != 0].shape[0] else 0).dropna().max()

实际上，我需要计算每个窗口的第一个非零值和最后一个值之间的差异。上面的代码有效，但我想知道是否有更有效的方法来执行相同的操作。

谢谢你的帮助。

标签： pythonpandas

解决方案

你的代码的问题是它对 0 执行了太多的比较。由于dist列中的值没有改变，所以比较一次并重用结果。

下面是这样做的方法：

# The scale of our problem
n = 200_000
window_size = 3_000

# Some simulated data. 20% of `dist` is zero
n_zero = int(n * 0.2)
dist = np.hstack([np.zeros(n_zero), np.random.uniform(0, 1, n - n_zero)])
np.random.shuffle(dist)

df = pd.DataFrame({
    'dist': dist
})

# -----------------

# Convert dist to a numpy array. We do not need pandas series here
dist = df['dist'].to_numpy()

# Find the indexes of all non-zero elements. nz = non-zero
nz_index, = np.nonzero(dist)

# For each row in `dist`, find the first non-zero value within
# the next `window_size` rows. Return NaN if no such value
# can be found.
dist_nz = np.empty_like(dist)
idx = 0
for i in range(len(dist)):
    idx += 1 if i > nz_index[idx] else 0
    dist_nz[i] = dist[nz_index[idx]] if idx < len(nz_index) and (idx - i < window_size) else np.nan

(dist[window_size-1:] - dist_nz[:-window_size+1]).max()

与原始的 2m 13s 相比，这在 0.3s 内完成。

推荐阅读