stock/Chan241206/find_strokes.py

def find_strokes(fractals, df, min_interval=5):
    """
    在此函数中对初步分型（fractals）进行有效性判断和筛选，
    包括：
    1. 去除连续同类型分型中较不极端的分型。
    2. 确保顶底交替和满足最小间隔要求 min_interval。
    3. 最终根据有效分型构造笔（strokes）。

    :param fractals: 初步的分型点列表 [(index, 'Top'), (index, 'Bottom')]
    :param df: 去除包含关系的K线数据，包含High、Low
    :param min_interval: 顶底之间的最小间隔K线数
    :return: (valid_fractals, strokes)
        valid_fractals: 有效分型列表 [(index, 'Top'), (index, 'Bottom')]
        strokes: 笔的列表 [(start_index, end_index)]
    """
    if not fractals:
        return [], []

    # 1. 先对分型列表排序（按index）
    fractals = sorted(fractals, key=lambda x: x[0])

    valid_fractals = [fractals[0]]  # 将第一个分型作为初始有效分型
    # 开始进行有效性筛选
    for i in range(1, len(fractals)):
        current_idx, current_type = fractals[i]
        last_idx, last_type = valid_fractals[-1]

        # 如果类型相同，保留更加极端的分型
        if current_type == last_type:
            if current_type == 'Top':
                # 保留更高的顶
                if df.loc[current_idx, 'High'] > df.loc[last_idx, 'High']:
                    valid_fractals[-1] = (current_idx, current_type)
                # 否则保持不变（跳过当前分型）
            else:  # 'Bottom'
                # 保留更低的底
                if df.loc[current_idx, 'Low'] < df.loc[last_idx, 'Low']:
                    valid_fractals[-1] = (current_idx, current_type)
                # 否则保持不变
        else:
            # 类型不同，检查间隔
            if current_idx - last_idx - 1 >= min_interval:
                valid_fractals.append((current_idx, current_type))
            # 如果间隔不够，不加入

    # 有了有效分型列表后，再构造笔
    strokes = []
    for i in range(len(valid_fractals) - 1):
        start_idx, start_type = valid_fractals[i]
        end_idx, end_type = valid_fractals[i + 1]
        # 此时有效分型保证类型交替（如果需要进一步确保，也可再加判断）
        strokes.append((start_idx, end_idx))

    return valid_fractals, strokes