Python中get_list_axes函数y轴列表生成异常及优化方案问询

需求说明 #

函数get_list_axes接收两个参数：

• 唯一captors列表（示例：["A", "B", "C", "D"]）
• 需共享同一y轴的captors子集（示例：["A", "D"]、["A", "C"]等）

函数目标：根据合并子集返回对应y轴列表，无合并需求时返回["y1", "y2", "y3", "y4"]。

当前实现代码 #

def get_list_axes(list_ana_captors: list, merge_captors: list) -> list:
    counter_no_merge, counter_merge = 1, 1
    list_axes = ["y1"]

    for captor in list_ana_captors[1:]:
        if captor not in merge_captors or captor == merge_captors[0]:
            if captor == merge_captors[0]:
                counter_merge = counter_no_merge + 1
            counter_no_merge = counter_no_merge + 1

        if captor not in merge_captors:
            list_axes.append(f"y{counter_no_merge}")
        else:
            list_axes.append(f"y{counter_merge}")
    return list_axes

测试情况 #

定义list_captors = ["A", "B", "C", "D"]：

符合预期的场景 #

• get_list_axes(list_captors, ["A", "B", "C", "D"]) = ['y1', 'y1', 'y1', 'y1']
• get_list_axes(list_captors, ["B", "C", "D"]) = ['y1', 'y2', 'y2', 'y2']
• get_list_axes(list_captors, ["A", "B"]) = ['y1', 'y1', 'y2', 'y3']
• get_list_axes(list_captors, ["A", "C"]) = ['y1', 'y2', 'y1', 'y3']
• get_list_axes(list_captors, ["A", "D"]) = ['y1', 'y2', 'y3', 'y1']
• get_list_axes(list_captors, ["B", "D"]) = ['y1', 'y2', 'y3', 'y2']

不符合预期的场景 #

• get_list_axes(list_captors, ["D", "A"]) 返回['y1', 'y2', 'y3', 'y4']，预期应为['y1', 'y2', 'y3', 'y1']
• get_list_axes(list_captors, ["D", "C"]) 返回['y1', 'y2', 'y1', 'y3']，预期应为['y1', 'y2', 'y3', 'y3']
• get_list_axes(list_captors, ["D", "B"]) 返回['y1', 'y1', 'y2', 'y3']，预期应为['y1', 'y2', 'y3', 'y2']
• get_list_axes(list_captors, ["D", "C", "A", "B"]) 返回['y1', 'y1', 'y1', 'y2']，预期应为['y1', 'y1', 'y1', 'y1']

优化实现方案 #

当前实现依赖merge_captors的顺序，导致顺序不符时出错。更简洁可靠的方式是先建立captor到y轴编号的映射：

1. 遍历主captors列表，找到合并组中第一个出现的元素，它的y轴编号就是整个共享组的编号
2. 再次遍历主列表，为每个元素分配编号：不在合并组的按顺序递增，在合并组的统一使用共享编号

优化后代码 #

def get_list_axes(list_ana_captors: list, merge_captors: list) -> list:
    merge_set = set(merge_captors)
    y_mapping = {}
    current_y = 1
    
    # 确定合并组的共享y编号（取合并组中第一个出现在主列表的元素的编号）
    shared_y = None
    for captor in list_ana_captors:
        if captor in merge_set:
            shared_y = current_y
            break
        current_y += 1
    
    # 重置计数器，构建完整的captor到y轴的映射
    current_y = 1
    for captor in list_ana_captors:
        if captor in merge_set:
            y_mapping[captor] = f"y{shared_y}"
        else:
            y_mapping[captor] = f"y{current_y}"
            current_y += 1
    
    # 按主列表顺序生成结果
    return [y_mapping[captor] for captor in list_ana_captors]

验证优化后结果 #

测试所有原不符合预期的场景，均返回正确结果：

• get_list_axes(list_captors, ["D", "A"]) → ['y1', 'y2', 'y3', 'y1']
• get_list_axes(list_captors, ["D", "C"]) → ['y1', 'y2', 'y3', 'y3']
• get_list_axes(list_captors, ["D", "B"]) → ['y1', 'y2', 'y3', 'y2']
• get_list_axes(list_captors, ["D", "C", "A", "B"]) → ['y1', 'y1', 'y1', 'y1']

同时原符合预期的场景也能正常返回正确结果。

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