如何将未知深度的决策树对象转换为Plotly.js Sankey图数据集？

我来帮你搞定这个问题——把任意深度的决策树转换成Sankey图需要的数据集其实不难，核心就是用递归遍历的方式，不管树有多深都能自动处理，同时维护好节点的唯一标识和连接关系就行。

首先先明确你提供的决策树结构（方便对照）：

"DecisionTreeRegressionModel": [
  {
    " If (feature 28 <= 16.0)": [
      {
        " If (feature 0 in {0.0})": [
          { " Predict: 13.0": [] }
        ]
      },
      {
        " Else (feature 0 not in {0.0})": [
          { " Predict: 16.0": [] }
        ]
      }
    ]
  },
  {
    " Else (feature 28 > 16.0)": [
      {
        " If (feature 28 <= 40.0)": [
          { " Predict: 40.0": [] }
        ]
      },
      {
        " Else (feature 28 > 40.0)": [
          {
            " If (feature 0 in {0.0})": [
              { " Predict: 45.0": [] }
            ]
          },
          {
            " Else (feature 0 not in {0.0})": [
              { " Predict: 50.0": [] }
            ]
          }
        ]
      }
    ]
  }
]

核心转换逻辑 #

Sankey图依赖两个核心数据集：

• nodes：所有节点的列表，每个节点需要唯一ID和显示标签
• links：节点间的连接列表，每个连接包含源节点ID、目标节点ID和流量值（这里决策树分支可以统一设为1）

递归遍历的关键步骤：

• 维护一个全局计数器，给每个新节点分配唯一ID
• 维护nodes和links数组存储最终数据
• 对每个节点：
  1. 提取节点的文本描述（比如"If (feature 28 <=16.0)"）
  2. 将节点加入nodes数组
  3. 如果有父节点，添加从父节点到当前节点的连接到links
  4. 如果当前节点有子节点，递归处理每个子节点

JavaScript实现示例（适配Plotly Sankey） #

// 全局状态：节点ID计数器、Sankey节点列表、连接列表
let nodeIdCounter = 0;
const sankeyNodes = [];
const sankeyLinks = [];

/**
 * 递归处理决策树节点
 * @param {number|null} parentId - 父节点ID，根节点传null
 * @param {object} nodeData - 当前要处理的决策树节点数据
 */
function traverseDecisionTree(parentId, nodeData) {
  // 遍历当前节点的每个键值对（每个节点是{节点描述: 子节点数组}的结构）
  for (const [nodeLabel, children] of Object.entries(nodeData)) {
    // 为当前节点分配唯一ID
    const currentId = nodeIdCounter++;
    // 添加到Sankey节点列表
    sankeyNodes.push({
      id: currentId,
      label: nodeLabel.trim()
    });
    // 如果有父节点，添加连接关系
    if (parentId !== null) {
      sankeyLinks.push({
        source: parentId,
        target: currentId,
        value: 1 // 决策树分支流量统一设为1，也可根据需求调整
      });
    }
    // 递归处理子节点（如果存在）
    if (Array.isArray(children) && children.length > 0) {
      children.forEach(childNode => {
        traverseDecisionTree(currentId, childNode);
      });
    }
  }
}

// 启动转换：传入你的决策树数据
const decisionTree = {
  "DecisionTreeRegressionModel": [
    // 这里放入你的完整决策树数据
  ]
};
traverseDecisionTree(null, decisionTree);

// 生成Plotly可直接使用的Sankey数据格式
const plotlySankeyData = {
  type: 'sankey',
  node: {
    pad: 15,
    thickness: 20,
    line: { color: "black", width: 0.5 },
    label: sankeyNodes.map(node => node.label)
  },
  link: {
    source: sankeyLinks.map(link => link.source),
    target: sankeyLinks.map(link => link.target),
    value: sankeyLinks.map(link => link.value)
  }
};

// 渲染示例：替换为你的容器ID
// Plotly.newPlot('sankey-chart', [plotlySankeyData]);

关键细节说明 #

• 递归适配任意深度：不管决策树嵌套多少层，递归函数都会逐层遍历，不会遗漏任何节点
• 唯一节点ID：用nodeIdCounter保证每个节点的ID唯一，避免Sankey图连接混乱
• 末端节点自动处理：预测节点（比如"Predict:13.0"）没有子节点，递归到这里会自动停止，成为Sankey图的末端节点
• 流量值灵活调整：如果需要区分分支权重，可以根据决策树的样本数量等信息修改value的值，这里统一设为1是因为决策树分支是互斥的

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