K DISTANCE

Given a balanced binary tree of integars and an integar B, count the number of pair of B distant nodes (a,b) where:

  1. a is ancestor of b
  2. value of node a - value of node b | <= B

Constraints

1 <= Number of nodes in binary tree <= 100000
0 <= node values <= 10^5
1 <= B <= 10^5 

For Example

Input 1:
            1
          /   \
         2    3
        / \  / \
       4   5 6  7
      /
     8 
     B = 1
Output 1:
    1
Explaination 1:
    Possible pairs are :- {(1, 2)}

Input 2:
            1
           /  \
          2    3
           \
            4
             \
              5
    B = 2
Output 2:
    4
Explaination 2:
Possible pairs are :- {(1, 2), (1, 3), (2, 4), (4, 5)}
Solution:

Time: O(nlogn)

/**
 * Definition for binary tree
 * class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) {
 *      val = x;
 *      left=null;
 *      right=null;
 *     }
 * }
 */
public class Solution {
    public int solve(TreeNode A, int B) {
        List<Integer> ancestors = new ArrayList<>();
        ancestors.add(A.val);
        return helper(A.left, ancestors, B) + helper(A.right, ancestors, B);
    }
    
    private int helper(TreeNode A, List<Integer> ancestors, int B) {
        if (A == null) {
            return 0;
        }
        ancestors.add(A.val);
        int left = helper(A.left, ancestors, B);
        int right = helper(A.right, ancestors, B);
        ancestors.remove(ancestors.size() - 1);

        int count = 0;
        for (int ancestor : ancestors) {
            if (Math.abs(A.val - ancestor) <= B) {
                count ++;
            }
        }
        return count + left + right;
    }
}