Sum of distances in tree

/*
 * There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.
 * You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.
 * Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.
 * Example 1:
 * Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
 * Output: [8,12,6,10,10,10]
 * Explanation: The tree is shown above.
 * We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
 * equals 1 + 1 + 2 + 2 + 2 = 8.
 * Hence, answer[0] = 8, and so on.
 * Example 2:
 * Input: n = 1, edges = []
 * Output: [0]
 * Example 3:
 * Input: n = 2, edges = [[1,0]]
 * Output: [1,1]
 * Constraints:
 * 1 <= n <= 3 * 104
 * edges.length == n - 1
 * edges[i].length == 2
 * 0 <= ai, bi < n
 * ai != bi
 * The given input represents a valid tree.
 */

impl Solution {
  pub fn sum_of_distances_in_tree(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
    let n = n as usize;
    let mut graph = vec![vec![]; n];
    for e in &edges {
      let (u, v) = (e[0] as usize, e[1] as usize);
      graph[u].push(v);
      graph[v].push(u);
    }
    let mut count = vec![1usize; n];
    let mut ans = vec![0i64; n];
    // First DFS: root at 0
    fn dfs1(v: usize, p: usize, graph: &Vec<Vec<usize>>, count: &mut Vec<usize>, ans: &mut Vec<i64>) {
      for &u in &graph[v] {
        if u != p {
          dfs1(u, v, graph, count, ans);
          count[v] += count[u];
          ans[v] += ans[u] + count[u] as i64;
        }
      }
    }
    // Second DFS: reroot
    fn dfs2(v: usize, p: usize, n: usize, graph: &Vec<Vec<usize>>, count: &mut Vec<usize>, ans: &mut Vec<i64>) {
      for &u in &graph[v] {
        if u != p {
          ans[u] = ans[v] - count[u] as i64 + (n - count[u]) as i64;
          dfs2(u, v, n, graph, count, ans);
        }
      }
    }
    dfs1(0, n, &graph, &mut count, &mut ans);
    dfs2(0, n, n, &graph, &mut count, &mut ans);
    ans.iter().map(|&x| x as i32).collect()
  }
}