LeetCode 1245. Tree Diameter

Problem Statement

use std::collections::HashMap;
use std::cmp;

impl Solution {
    fn dfs(g: &HashMap<i32, Vec<i32>>, ans: &mut i32, cur: i32, parent: i32) -> i32 {
        let val = g.get(&cur).unwrap();
        if val.len() == 1 && val[0] == parent {
            return 1;
        }

        let (mut a, mut b) = (0, 0);
        for item in val {
            if *item == parent {
                continue;
            }
            let next = Self::dfs(g, ans, *item, cur);
            if next > a {
                b = a;
                a = next;
            } else if next > b {
                b = next;
            }
            *ans = cmp::max(*ans, a + b);
        }
        return a + 1;
    }

    pub fn tree_diameter(edges: Vec<Vec<i32>>) -> i32 {
        let (mut tree, mut ans) : (HashMap<i32, Vec<i32>>, i32) = (HashMap::new(), 0);
        for edge in &edges {
            let (k, v) = (edge[0], edge[1]);
            tree.entry(k).or_insert(vec![]);
            tree.entry(v).or_insert(vec![]);
            tree.get_mut(&k).unwrap().push(v);
            tree.get_mut(&v).unwrap().push(k);
        }
        Self::dfs(&tree, &mut ans, edges[0][0], -1);
        return ans;
    }
}