如何在D3.js Tidy Tree图中查找选中节点的关联节点

问题背景 #

我现有一个基于D3.js实现的Tidy Tree图表，代码如下：

JavaScript代码 #

const data = {
  "name": "tree",
  "children": [{
    "name": "A",
    "children": [{
        "name": "B",
        "children": [{
          "name": "C",
          "size": 3938
        }, ]
      },
      {
        "name": "H",
        "children": [{
            "name": "D",
            "size": 3534
          },
          {
            "name": "E",
            "size": 5731
          },
          {
            "name": "F",
            "size": 7840
          },
        ]
      },
    ]
  }]
}

const width = 928;

// Compute the tree height; this approach will allow the height of the
// SVG to scale according to the breadth (width) of the tree layout.
const root = d3.hierarchy(data);
const dx = 10;
const dy = width / (root.height + 1);

// Create a tree layout.
const tree = d3.tree().nodeSize([dx, dy]);

// Sort the tree and apply the layout.
root.sort((a, b) => d3.ascending(a.data.name, b.data.name));
tree(root);

// Compute the extent of the tree. Note that x and y are swapped here
// because in the tree layout, x is the breadth, but when displayed, the
// tree extends right rather than down.
let x0 = Infinity;
let x1 = -x0;
root.each(d => {
  if (d.x > x1) x1 = d.x;
  if (d.x < x0) x0 = d.x;
});

// Compute the adjusted height of the tree.
const height = x1 - x0 + dx * 2;

const svg = d3.select("svg")
  .attr("width", width)
  .attr("height", height)
  .attr("viewBox", [-dy / 3, x0 - dx, width, height])
  .attr("style", "max-width: 100%; height: auto; font: 10px sans-serif;");

const link = svg.append("g")
  .attr("fill", "none")
  .attr("stroke", "#555")
  .attr("stroke-opacity", 0.4)
  .attr("stroke-width", 1.5)

link
  .selectAll()
  .data(root.links())
  .join("path")
  .attr("d", d3.linkHorizontal()
    .x(d => d.y)
    .y(d => d.x));

const node = svg.append("g")
  .attr("stroke-linejoin", "round")
  .attr("stroke-width", 3)
  .selectAll()
  .data(root.descendants())
  .join("g")
  .attr("transform", d => `translate(${d.y},${d.x})`)

node.append("circle")
  .attr("fill", d => d.children ? "#555" : "#999")
  .attr("r", 5.5)

node.append("text")
  .attr("dy", "0.31em")
  .attr("x", d => d.children ? -6 : 6)
  .attr("text-anchor", d => d.children ? "end" : "start")
  .text(d => d.data.name)
  .clone(true).lower()
  .attr("stroke", "white");

d3.selectAll("circle").on("contextmenu", (event) => {
  this.selectedNode = event.srcElement.__data__.data
});

const connectNodes = (t, f) => {
  let to = null,
    fr = null;
  node.each(d => {
    if (d.data.name === t) to = d;
    if (d.data.name === f) fr = d;
  });
  if (to && fr) {
    link.append("path")
      .attr("d", "M" + to.y + "," + to.x + "L" + fr.y + "," + fr.x)
      .attr("fill", "none")
      .attr("stroke", "red");
  }
};
connectNodes("B", "D");

HTML代码 #

<!doctype html>
<html>

<head>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/d3/7.8.5/d3.js"></script>
</head>

<body>
  <svg></svg>
</body>

</html>

需求：当右键选中某个节点（例如节点B）时，获取该节点的所有关联节点——包括父节点（A）、子节点（C）以及通过connectNodes添加的自定义连接节点（D）。

解决方案 #

要实现这个功能，需要分三步处理：维护自定义连接关系、编写关联节点查询函数、整合到选中事件中。

1. 维护自定义连接关系 #

原代码中connectNodes仅绘制连线，未存储连接关系。我们需要新增一个集合来记录所有自定义连接，方便后续查询：

// 新增：存储所有自定义双向连接（避免重复）
const customConnections = new Set();

// 修改connectNodes函数，同时记录连接关系
const connectNodes = (t, f) => {
  let to = null,
    fr = null;
  node.each(d => {
    if (d.data.name === t) to = d;
    if (d.data.name === f) fr = d;
  });
  if (to && fr) {
    link.append("path")
      .attr("d", "M" + to.y + "," + to.x + "L" + fr.y + "," + fr.x)
      .attr("fill", "none")
      .attr("stroke", "red");
    // 用排序后的字符串存储，确保"A-B"和"B-A"视为同一连接
    const pair = [t, f].sort().join("-");
    customConnections.add(pair);
  }
};
connectNodes("B", "D");

2. 编写获取关联节点的函数 #

该函数需要遍历目标节点的父节点、子节点，同时查询自定义连接中的关联节点：

// 获取指定节点的所有关联节点
const getRelatedNodes = (targetNodeName) => {
  const related = new Set();
  let targetNode = null;

  // 先找到目标节点的d3.hierarchy对象
  root.each(d => {
    if (d.data.name === targetNodeName) {
      targetNode = d;
    }
  });

  if (!targetNode) return [];

  // 1. 添加父节点
  if (targetNode.parent) {
    related.add(targetNode.parent.data.name);
  }

  // 2. 添加子节点
  if (targetNode.children) {
    targetNode.children.forEach(child => {
      related.add(child.data.name);
    });
  }

  // 3. 添加自定义连接的节点
  customConnections.forEach(pair => {
    const [a, b] = pair.split("-");
    if (a === targetNodeName) {
      related.add(b);
    } else if (b === targetNodeName) {
      related.add(a);
    }
  });

  // 转换为数组返回，排除自身
  return Array.from(related).filter(name => name !== targetNodeName);
};

3. 整合到右键选中事件中 #

修改右键菜单事件，调用getRelatedNodes函数获取并处理关联节点：

d3.selectAll("circle").on("contextmenu", (event) => {
  event.preventDefault(); // 阻止默认右键菜单弹出
  const selectedData = event.srcElement.__data__.data;
  const relatedNodes = getRelatedNodes(selectedData.name);
  
  // 控制台输出结果（可根据需求替换为其他逻辑）
  console.log(`选中节点：${selectedData.name}`);
  console.log("关联节点：", relatedNodes);
  
  // 可选：高亮选中节点和关联节点
  node.selectAll("circle").attr("fill", d => {
    if (relatedNodes.includes(d.data.name) || d.data.name === selectedData.name) {
      return "#ff0000"; // 红色高亮
    }
    return d.children ? "#555" : "#999"; // 恢复默认颜色
  });
});

效果验证 #

右键点击节点B时，控制台会输出：

选中节点：B
关联节点： ["A", "C", "D"]

同时图表中B、A、C、D节点会被红色高亮，完全符合需求。

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