d3 / d3-hierarchy

2D layout algorithms for visualizing hierarchical data.
https://d3js.org/d3-hierarchy
ISC License
1.13k stars 315 forks source link

New nodesAtDepth(depth) function #178

Open benjaminpreiss opened 3 years ago

benjaminpreiss commented 3 years ago

Hey Mike

Great work. For tree traversal it would be great to have a function that returns all nodes at a certain depth. Example (taken from source code for leaves()):

export default function(depth) {
  var nodes = [];
  this.eachBefore(function(node) {
    if (node.depth === depth) {
      nodes.push(node);
    }
  });
  return nodes;
}

kind regards Benjamin

Fil commented 3 years ago

We should probably cut short when we're at the desired depth, and avoid a traversal of the branches that are deeper.

benjaminpreiss commented 3 years ago

@Fil Yeah, true... Thinking about this, it would be useful to extend iterator.js as well as create another function called eachToDepth:

First, the changes to iterator.js:

export default function*(depth = this.height) {
  var node = this, current, next = [node], children, i, n;
  do {
    current = next.reverse(), next = [];
    while (node = current.pop()) {
      yield node;
      if ((children = node.children) && node.depth < depth) {
        for (i = 0, n = children.length; i < n; ++i) {
          next.push(children[i]);
        }
      }
    }
  } while (next.length);
}

Now the new eachToDepth.js file (leveraging our changed iterator):

export default function(callback, depth, that) {
  let index = -1;
  for (const node of this[Symbol.iterator](depth)) {
    callback.call(that, node, ++index, this);
  }
  return this;
}

Then we can use the changes above to construct nodesAtDepth.js:

export default function(depth) {
  const nodes = []
  this.eachToDepth(function(node) {
    if (node.depth === depth) {
      nodes.push(node);
    }
  }, depth);
  return nodes
}

The above proposals include nodes at the specified depth (traverse tree breadth first including descendants with specified depth)

benjaminpreiss commented 3 years ago

Or maybe we don't even need eachToDepth.js...

In that case nodesAtDepth.js would be:

export default function(depth) {
  const nodes = []
  for (const node of this[Symbol.iterator](depth)) {
    if (node.depth === depth) {
      nodes.push(node);
    }
  }
  return nodes
}