Number of Good Leaf Nodes Pairs LeetCode solution in JavaScript

1
/**
2
 * Definition for a binary tree node.
3
 * function TreeNode(val, left, right) {
4
 *     this.val = (val===undefined ? 0 : val)
5
 *     this.left = (left===undefined ? null : left)
6
 *     this.right = (right===undefined ? null : right)
7
 * }
8
 */
9
/**
10
 * @param {TreeNode} root
11
 * @param {number} distance
12
 * @return {number}
13
 */
14
var countPairs = function (root, distance) {
15
  var graph = new Map();
16
  var leaves = new Map();
17
  function dfs(root, parent) {
18
    if (!root) return;
19
    if (!root.left && !root.right) {
20
      leaves.set(root, true);
21
    }
22

23
    graph.set(root, []);
24
    if (root.left) {
25
      graph.get(root).push(root.left);
26
    }
27
    if (root.right) {
28
      graph.get(root).push(root.right);
29
    }
30
    if (parent) {
31
      graph.get(root).push(parent);
32
    }
33
    dfs(root.left, root);
34
    dfs(root.right, root);
35
  }
36

37
  dfs(root, null);
38

39
  var visited = new Map();
40
  var count = 0;
41

42
  function bfs(start, dis) {
43
    visited.set(start, true);
44
    var queue = [graph.get(start).filter((node) => !visited.has(node))];
45

46
    var innerVisited = new Map();
47

48
    while (queue.length) {
49
      var curLevelNodes = queue.shift();
50

51
      if (!curLevelNodes.length) break;
52
      if (dis === 0) break;
53

54
      var nextLevelNodes = [];
55
      for (var i = 0; i < curLevelNodes.length; i++) {
56
        var curLevelNode = curLevelNodes[i];
57

58
        innerVisited.set(curLevelNode, true);
59

60
        if (leaves.has(curLevelNode)) {
61
          count++;
62
        }
63

64
        nextLevelNodes.push(
65
          ...graph
66
            .get(curLevelNode)
67
            .filter((node) => !visited.has(node) && !innerVisited.has(node))
68
        );
69
      }
70
      queue.push(nextLevelNodes);
71
      dis--;
72
    }
73
  }
74

75
  leaves.forEach((value, leaf) => {
76
    bfs(leaf, distance);
77
  });
78

79
  return count;
80
};