Problem Statement
Medium
You are given the root of a binary tree and an integer distance. A pair of two different leaf nodes of a binary tree is said to be good if the length of the shortest path between them is less than or equal to distance.
Return the number of good leaf node pairs in the tree.
Example 1:
Input: root = [1,2,3,null,4], distance = 3
Output: 1
Explanation: The leaf nodes of the tree are 3 and 4 and the length of the shortest path between them is 3. This is the only good pair.
Example 2:
Input: root = [1,2,3,4,5,6,7], distance = 3
Output: 2
Explanation: The good pairs are [4,5] and [6,7] with shortest path = 2. The pair [4,6] is not good because the length of ther shortest path between them is 4.
Example 3:
Input: root = [7,1,4,6,null,5,3,null,null,null,null,null,2], distance = 3
Output: 1
Explanation: The only good pair is [2,5].
Constraints:
- The number of nodes in the
tree is in the range [1, 210].
1 <= Node.val <= 1001 <= distance <= 10
/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode() {}
* TreeNode(int val) { this.val = val; }
* TreeNode(int val, TreeNode left, TreeNode right) {
* this.val = val;
* this.left = left;
* this.right = right;
* }
* }
*/
class Solution {
Map<TreeNode, List<TreeNode>> adjList = new HashMap<>();
Set<TreeNode> leaves = new HashSet<>();
private TreeNode dfs(TreeNode root) {
if (root == null) {
return null;
}
if (root.left == null && root.right == null) {
leaves.add(root);
}
if (root.left != null) {
adjList.computeIfAbsent(root, k -> new ArrayList<>()).add(dfs(root.left));
adjList.computeIfAbsent(root.left, k -> new ArrayList<>()).add(root);
}
if (root.right != null) {
adjList.computeIfAbsent(root, k -> new ArrayList<>()).add(dfs(root.right));
adjList.computeIfAbsent(root.right, k -> new ArrayList<>()).add(root);
}
return root;
}
public int countPairs(TreeNode root, int distance) {
dfs(root); // Generate Adjacency List & leaf nodes
int ans = 0;
for(TreeNode leaf: leaves){ // shortest path algorithm for each destination (leaf)
Queue<TreeNode> bfs = new LinkedList<>(); // to track 1+ hop neighbors of leaf
Set<TreeNode> visited = new HashSet<>();
bfs.add(leaf);
visited.add(leaf);
for(int i=0; i<=distance; i++){
int n_iHopNeighbors = bfs.size();
for(int j = 0; j<n_iHopNeighbors; j++){
TreeNode curr = bfs.poll();
if(leaves.contains(curr) && curr!=leaf)
ans++;
for(TreeNode neighborOfCurr: adjList.getOrDefault(curr, new ArrayList<>())){
if(!visited.contains(neighborOfCurr)){
bfs.add(neighborOfCurr);
visited.add(neighborOfCurr);
}
}
}
}
}
return ans/2;
}
}