Python 数组体积大，时间成本高

您的原始代码给了我一个错误，因为数组中前几个条目的维度不匹配，这些条目太短了，因为窗口尚未满，因此我将其更改为丢弃第一个值：

def rolling_approach(df, window_size=3):
    return np.array(
        [w.values for w in df[["close", "volume"]].rolling(window=window_size)][
            window_size - 1 :
        ]
    )

PD的。对于这类操作，DataFrame.rolling 可能非常慢。shift 非常有效 下面是一个写出的例子：pandaswindow_size=3

pd.concat(
            [
                df[["close", "volume"]].shift().shift(),
                df[["close", "volume"]].shift(),
                df[["close", "volume"]],
            ],
            axis=1,
        )
        .values[2:, :]
        .reshape(-1, 3, 2)
    )

我们堆叠班次，然后重塑值。

推广到一个变量，我们得到：window_size

def shift_approach(df, window_size=3):
    shifted_df = pd.concat(
        [df[["close", "volume"]].shift(i) for i in range(window_size - 1, -1, -1)],
        axis=1,
    )
    reshaped_array = shifted_df.values[window_size - 1 :, :].reshape(-1, window_size, 2)
    return reshaped_array

shift性能比操作高出近两个数量级：rolling

在我的MacBook上，它可以很好地扩展到数亿行：

def setup(N):
    np.random.seed(42)
    close_values = np.random.randint(1, 100, size=N)
    volume_values = np.random.randint(100, 1000, size=N)
    df = pd.DataFrame({"close": close_values, "volume": volume_values})
    return [df, 10]


approaches = [rolling_approach, shift_approach]
# Show correctness
for approach in approaches[1:]:
    data = setup(10)
    assert np.isclose(approach(*data), approaches[0](*data)).all()

run_performance_comparison(
    approaches,
    [
        1000,
        3000,
        5000,
        10000,
        30000,
        100_000,
        300_000,
        500_000,
        1_000_000,
        3_000_000,
        5_000_000,
        10_000_000,
    ],
    setup=setup,
    title="Performance Comparison",
    number_of_repetitions=2,
)

分析代码：

import timeit
from functools import partial

import matplotlib.pyplot as plt
from typing import List, Dict, Callable

from contextlib import contextmanager


@contextmanager
def data_provider(data_size, setup=lambda N: N, teardown=lambda: None):
    data = setup(data_size)
    yield data
    teardown()


def run_performance_comparison(approaches: List[Callable],
                               data_size: List[int],
                               setup=lambda N: N,
                               teardown=lambda: None,
                               number_of_repetitions=5, title='Performance Comparison', data_name='N'):
    approach_times: Dict[Callable, List[float]] = {approach: [] for approach in approaches}
    for N in data_size:
        with data_provider(N, setup, teardown) as data:
            for approach in approaches:
                function = partial(approach, *data)
                approach_time = min(timeit.Timer(function).repeat(repeat=number_of_repetitions, number=2))
                approach_times[approach].append(approach_time)

    for approach in approaches:
        plt.plot(data_size, approach_times[approach], label=approach.__name__)
    plt.yscale('log')
    plt.xscale('log')

    plt.xlabel(data_name)
    plt.ylabel('Execution Time (seconds)')
    plt.title(title)
    plt.legend()
    plt.show()