Leetcode

Find Distance in a Binary Tree

May 26, 2022 by aAshish

  • Time:O(h)
  • Space:O(h)

C++

 
class Solution {
 public:
  int findDistance(TreeNode* root, int p, int q) {
    TreeNode* lca = getLCA(root, p, q);
    return dist(lca, p) + dist(lca, q);
  }

 private:
  TreeNode* getLCA(TreeNode* root, int p, int q) {
    if (!root || root->val == p || root->val == q)
      return root;

    TreeNode* l = getLCA(root->left, p, q);
    TreeNode* r = getLCA(root->right, p, q);

    if (l && r)
      return root;
    return l ? l : r;
  }

  int dist(TreeNode* lca, int target) {
    if (!lca)
      return 10000;
    if (lca->val == target)
      return 0;
    return 1 + min(dist(lca->left, target), dist(lca->right, target));
  }
};

JAVA

 
class Solution {
  public int findDistance(TreeNode root, int p, int q) {
    TreeNode lca = getLCA(root, p, q);
    return dist(lca, p) + dist(lca, q);
  }

  private TreeNode getLCA(TreeNode root, int p, int q) {
    if (root == null || root.val == p || root.val == q)
      return root;

    TreeNode l = getLCA(root.left, p, q);
    TreeNode r = getLCA(root.right, p, q);

    if (l != null && r != null)
      return root;
    return l == null ? r : l;
  }

  private int dist(TreeNode lca, int target) {
    if (lca == null)
      return 10000;
    if (lca.val == target)
      return 0;
    return 1 + Math.min(dist(lca.left, target), dist(lca.right, target));
  }
}

Python

 
class Solution:
  def findDistance(self, root: TreeNode, p: int, q: int) -> int:
    def getLCA(root, p, q):
      if not root or root.val == p or root.val == q:
        return root

      l = getLCA(root.left, p, q)
      r = getLCA(root.right, p, q)

      if l and r:
        return root
      return l or r

    def dist(lca, target):
      if not lca:
        return 10000
      if lca.val == target:
        return 0
      return 1 + min(dist(lca.left, target), dist(lca.right, target))

    lca = getLCA(root, p, q)
    return dist(lca, p) + dist(lca, q)