【Quant102】 经典技术指标 Pandas 实现（第一部分）（3）

def rsi_strategy(df, rsi_col='rsi', inplace=False):
    """
    实现RSI策略，生成买卖信号。
    
    参数:
    df : DataFrame
        包含RSI指标数据的DataFrame，必须包含'rsi_col'列。
    rsi_col : str
        RSI指标的列名。
    inplace : bool, optional
        是否在原地更新DataFrame，默认为False。
    
    返回:
    df : DataFrame
        原始DataFrame，增加了一个名为'signal'的列，包含买卖信号。
    """
    
    # 复制DataFrame以避免修改原始数据
    if not inplace:
        df = df.copy()
    
    # 初始化信号列，默认为无操作
    df['signal'] = 0
    
    # 产生买入信号：RSI低于30
    df.loc[df[rsi_col] < 30, 'signal'] = 1
    
    # 产生卖出信号：RSI高于70
    df.loc[df[rsi_col] > 70, 'signal'] = -1
    
    return df

def kdj(df, k_col='K', d_col='D', j_col='J', n=9, m1=3, m2=3, inplace=False):
    """
    计算KDJ指标，并更新DataFrame。
    
    参数:
    df : DataFrame
        包含价格数据的DataFrame，必须包含'open', 'high', 'low', 'close'列。
    n : int, optional
        计算KDJ指标的时间窗口，默认为9。
    m1 : int, optional
        计算D线的时间窗口，默认为3。
    m2 : int, optional
        计算J线的参数，默认为3。
    inplace : bool, optional
        是否在原地更新DataFrame，默认为False。
    
    返回:
    df : DataFrame
        包含KDJ指标的新DataFrame（如果inplace为False）。
    """
    
    # 复制DataFrame以避免修改原始数据
    if not inplace:
        df = df.copy()
    
    # 计算最小值和最大值
    low_min = df['low'].rolling(window=n).min()
    high_max = df['high'].rolling(window=n).max()
    
    # 计算K值
    df[k_col] = (df['close'] - low_min) / (high_max - low_min) * 100
    
    # 计算D值
    df[d_col] = df[k_col].rolling(window=m1).mean()
    
    # 计算J值
    df[j_col] = m2 * df[d_col] - df[k_col].rolling(window=m2).mean()
    
    return df