如何编写接收DataFrame的Python函数获取决策树所有可达路径

需求说明 #

• 输入：包含Node_ID、IF_TRUE、IF_FALSE三个字段的DataFrame，所有节点为二元分支，终节点固定为End_approve或End_deny
• 输出：所有从根节点（默认Node_ID=1）出发到终节点的有序路径列表，单条路径格式为[节点ID1, 节点ID2, …, 决策值]，决策值统一为小写的approve/deny

示例参考 #

输入样例 #

Node_ID IF_TRUE IF_FALSE
1 2 3
2 End_approve 3
3 End_deny End_approve

期望输出 #

[[1, 2, 'approve'], [1, 2, 3, 'deny'], [1, 2, 3, 'approve'], [1, 3, 'deny'], [1, 3, 'approve']]

实现代码 #

原参考代码为固定层数硬编码实现，仅支持深度不超过3层的决策树，以下为通用实现，可适配任意深度的二元决策树：

import pandas as pd
from collections import deque

def generate_decision_paths(df_in):
    # 构建节点映射字典，快速查询每个节点的真假分支
    node_map = df_in.set_index('Node_ID')[['IF_TRUE', 'IF_FALSE']].to_dict('index')
    result = []
    # 广度优先遍历队列，存储（当前节点ID, 当前路径列表）
    traverse_queue = deque()
    # 根节点默认为ID=1
    traverse_queue.append((1, [1]))
    
    while traverse_queue:
        current_id, current_path = traverse_queue.popleft()
        true_branch = node_map[current_id]['IF_TRUE']
        false_branch = node_map[current_id]['IF_FALSE']
        
        # 处理真分支
        if str(true_branch).startswith('End_'):
            # 到达终节点，转换决策值格式后加入结果
            decision = true_branch.split('_')[1].lower()
            result.append(current_path + [decision])
        else:
            # 中间节点，加入队列继续遍历
            traverse_queue.append((int(true_branch), current_path + [int(true_branch)]))
        
        # 处理假分支
        if str(false_branch).startswith('End_'):
            decision = false_branch.split('_')[1].lower()
            result.append(current_path + [decision])
        else:
            traverse_queue.append((int(false_branch), current_path + [int(false_branch)]))
    
    return result

# 测试代码
if __name__ == '__main__':
    test_df = pd.DataFrame({
        'Node_ID': [1, 2, 3],
        'IF_TRUE': [2, 'End_approve', 'End_deny'],
        'IF_FALSE': [3, 3, 'End_approve']
    })
    print(generate_decision_paths(test_df))

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