1740. Find Distance in a Binary Tree

Given the root of a binary tree and two integers p and q, return the distance between the nodes of value p and value q in the tree.

The distance between two nodes is the number of edges on the path from one to the other.

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 0
Output: 3
Explanation: There are 3 edges between 5 and 0: 5-3-1-0.

Example 2:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 7
Output: 2
Explanation: There are 2 edges between 5 and 7: 5-2-7.

Example 3:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 5
Output: 0
Explanation: The distance between a node and itself is 0.

 

Constraints:

• The number of nodes in the tree is in the range [1, 104].
• 0 <= Node.val <= 109
• All Node.val are unique.
• p and q are values in the tree.

Rust Solution

struct Solution;

use rustgym_util::*;
use std::collections::HashMap;
use std::collections::HashSet;

trait Preorder {
    fn preorder(&self, prev: i32, adj: &mut HashMap<i32, HashSet<i32>>);
}

impl Preorder for TreeLink {
    fn preorder(&self, prev: i32, adj: &mut HashMap<i32, HashSet<i32>>) {
        if let Some(node) = self {
            let node = node.borrow();
            let val = node.val;
            if prev != -1 {
                adj.entry(prev).or_default().insert(val);
                adj.entry(val).or_default().insert(prev);
            }
            node.left.preorder(val, adj);
            node.right.preorder(val, adj);
        }
    }
}

impl Solution {
    fn find_distance(root: TreeLink, p: i32, q: i32) -> i32 {
        let mut adj: HashMap<i32, HashSet<i32>> = HashMap::new();
        root.preorder(-1, &mut adj);
        Self::dfs(p, -1, &adj, q)
    }
    fn dfs(cur: i32, prev: i32, adj: &HashMap<i32, HashSet<i32>>, q: i32) -> i32 {
        if cur == q {
            0
        } else {
            for &next in adj[&cur].iter() {
                if next != prev {
                    let dist = Self::dfs(next, cur, adj, q);
                    if dist != -1 {
                        return dist + 1;
                    }
                }
            }
            -1
        }
    }
}

#[test]
fn test() {
    let root = tree!(
        3,
        tree!(5, tree!(6), tree!(2, tree!(7), tree!(4))),
        tree!(1, tree!(0), tree!(8))
    );
    let p = 5;
    let q = 0;
    let res = 3;
    assert_eq!(Solution::find_distance(root, p, q), res);
    let root = tree!(
        3,
        tree!(5, tree!(6), tree!(2, tree!(7), tree!(4))),
        tree!(1, tree!(0), tree!(8))
    );
    let p = 5;
    let q = 7;
    let res = 2;
    assert_eq!(Solution::find_distance(root, p, q), res);
    let root = tree!(
        3,
        tree!(5, tree!(6), tree!(2, tree!(7), tree!(4))),
        tree!(1, tree!(0), tree!(8))
    );
    let p = 5;
    let q = 5;
    let res = 0;
    assert_eq!(Solution::find_distance(root, p, q), res);
}