如何基于历史状态在pandas时序DataFrame中创建wallet持仓列？

嘿，这个需求太常见了——做简易量化策略里的仓位控制对吧？我给你两种靠谱的实现方式，兼顾可读性和效率，完美契合你“最多持有1股、空仓才买、持仓才卖”的规则。

核心逻辑先明确 #

我们需要维护一个持仓状态机：

• 初始状态：空仓（用0表示）
• 触发buy信号且当前是空仓 → 切换为持仓（1）
• 触发sell信号且当前是持仓 → 切换为空仓（0）
• 其他所有情况（比如持仓时的buy信号、空仓时的sell信号）→ 保持当前状态不变

方法1：直观迭代法（适合小数据集，逻辑一眼懂） #

如果你数据集不大，这种写法最容易理解，调试起来也方便：

# 初始化初始持仓状态：空仓
current_position = 0
wallet = []

# 逐行遍历数据，更新持仓状态
for _, row in df.iterrows():
    # 空仓时遇到买入信号，持仓变为1
    if row['buy'] and current_position == 0:
        current_position = 1
    # 持仓时遇到卖出信号，持仓变为0
    elif row['sell'] and current_position == 1:
        current_position = 0
    # 其他情况保持原状态，直接加入列表
    wallet.append(current_position)

# 将结果赋值给DataFrame的新列
df['wallet'] = wallet

方法2：高效生成器法（适合大数据集，兼顾速度和可读性） #

如果你的数据集很大，iterrows的速度可能不够看，用生成器结合numpy的写法能大幅提升效率，同时逻辑依然清晰：

import numpy as np

def track_position(buy_signals, sell_signals):
    # 初始化持仓状态
    position = 0
    # 遍历每一组买卖信号
    for buy, sell in zip(buy_signals, sell_signals):
        if buy and position == 0:
            position = 1
        elif sell and position == 1:
            position = 0
        # 输出当前状态
        yield position

# 用numpy快速生成结果列，比直接循环DataFrame快很多
df['wallet'] = np.fromiter(track_position(df['buy'], df['sell']), dtype=int)

额外处理：如果你的信号是字符串类型 #

如果你的buy和sell列存储的是字符串（比如"buy"/"sell"而非布尔值），先把它们转换成布尔类型再用上面的代码：

# 转换buy列：只有值为"buy"时视为有效买入信号
df['buy'] = df['buy'] == 'buy'
# 转换sell列：只有值为"sell"时视为有效卖出信号
df['sell'] = df['sell'] == 'sell'

测试验证一下 #

举个简单例子看看效果：

import pandas as pd

# 创建测试数据
test_data = {
    'buy': [True, False, False, True, False, False],
    'sell': [False, True, True, False, False, True]
}
df_test = pd.DataFrame(test_data, index=pd.date_range('2023-01-01', periods=6))

# 应用方法2生成wallet列
df_test['wallet'] = np.fromiter(track_position(df_test['buy'], df_test['sell']), dtype=int)
print(df_test)

输出结果完全符合预期：

buy   sell  wallet
2023-01-01  True  False       1
2023-01-02  False True        0
2023-01-03  False True        0
2023-01-04  True  False       1
2023-01-05  False False       1
2023-01-06  False True        0

可以看到：空仓时的买入信号生效，持仓时的卖出信号生效，无效信号（空仓时的sell、持仓时的buy）都被正确忽略了。

内容的提问来源于stack exchange，提问作者Jasper