DP on Trees Diameter of Tree in DSA Rust

Mental Model

The diameter of a tree is the longest path between any two nodes. We find it by checking the longest paths through each node using dynamic programming.

Analogy: Imagine a network of roads connecting towns. The diameter is the longest route you can drive without repeating roads. We find it by looking at the longest two roads leaving each town and combining them.

    1
   / \
  2   3
 /     \
4       5

Each number is a town (node), lines are roads (edges). Diameter is the longest path, e.g., 4 -> 2 -> 1 -> 3 -> 5.

Dry Run Walkthrough

Input: Tree edges: 1-2, 1-3, 2-4, 3-5

Goal: Find the longest path between any two nodes in the tree (diameter).

Step 1: Start DFS from node 1, calculate longest path down each child.

At node 1: children 2 and 3 to explore.

Why: We need to find longest paths from children to combine for diameter.

Step 2: At node 2, explore child 4. Node 4 has no children, so longest path down is 0.

Node 4: longest_down = 0

Why: Leaf node has no further paths.

Step 3: Back to node 2, longest path down is 1 (edge to 4).

Node 2: longest_down = 1 (via 4)

Why: Longest path down from 2 is through 4.

Step 4: At node 3, explore child 5. Node 5 is leaf, longest_down = 0.

Node 5: longest_down = 0

Why: Leaf node has no further paths.

Step 5: Back to node 3, longest path down is 1 (edge to 5).

Node 3: longest_down = 1 (via 5)

Why: Longest path down from 3 is through 5.

Step 6: At node 1, combine two longest child paths: 1 (from 2) + 1 (from 3) + 2 edges between them.

Diameter candidate at node 1 = 1 + 1 = 2 edges (path 4 -> 2 -> 1 -> 3 -> 5)

Why: Diameter passes through node 1 combining two longest child paths.

Step 7: Compare diameter candidates from all nodes, max is 2 edges.

Diameter = 2 edges, path length = 3 nodes

Why: Longest path found is the diameter.

Result:

4 -> 2 -> 1 -> 3 -> 5 -> null
Diameter length: 2 edges (3 nodes)

Annotated Code

DSA Rust

use std::collections::HashMap;

struct Tree {
    adj: HashMap<usize, Vec<usize>>,
    diameter: usize,
}

impl Tree {
    fn new() -> Self {
        Tree {
            adj: HashMap::new(),
            diameter: 0,
        }
    }

    fn add_edge(&mut self, u: usize, v: usize) {
        self.adj.entry(u).or_default().push(v);
        self.adj.entry(v).or_default().push(u);
    }

    fn dfs(&mut self, node: usize, parent: usize) -> usize {
        let mut max1 = 0; // longest path down
        let mut max2 = 0; // second longest path down

        for &child in self.adj.get(&node).unwrap_or(&vec![]) {
            if child == parent {
                continue; // avoid going back to parent
            }
            let length = self.dfs(child, node) + 1; // path length down child

            if length > max1 {
                max2 = max1;
                max1 = length;
            } else if length > max2 {
                max2 = length;
            }
        }

        self.diameter = self.diameter.max(max1 + max2); // update diameter
        max1
    }
}

fn main() {
    let mut tree = Tree::new();
    tree.add_edge(1, 2);
    tree.add_edge(1, 3);
    tree.add_edge(2, 4);
    tree.add_edge(3, 5);

    tree.dfs(1, 0);
    println!("Diameter length: {} edges", tree.diameter);
}

for &child in self.adj.get(&node).unwrap_or(&vec![]) {

iterate over children of current node

if child == parent { continue; }

skip parent to avoid revisiting

let length = self.dfs(child, node) + 1;

recursive call to get longest path down child plus edge

if length > max1 { max2 = max1; max1 = length; } else if length > max2 { max2 = length; }

track top two longest child paths

self.diameter = self.diameter.max(max1 + max2);

update diameter with sum of two longest paths through node

max1

return longest path down from current node

Output

Diameter length: 4 edges

Complexity Analysis

Time: O(n) because we visit each node once in DFS

Space: O(n) for adjacency list and recursion stack

vs Alternative: Naive approach checking all pairs is O(n^2), DP on trees is much faster

Edge Cases

Single node tree

Diameter is 0 edges since no paths exist

DSA Rust

for &child in self.adj.get(&node).unwrap_or(&vec![]) {

Linear chain tree (like a linked list)

Diameter equals number of edges in chain

DSA Rust

self.diameter = self.diameter.max(max1 + max2);

Tree with multiple branches of same length

Diameter correctly picks longest combined path

DSA Rust

if length > max1 { max2 = max1; max1 = length; } else if length > max2 { max2 = length; }

When to Use This Pattern

When asked for longest path between any two nodes in a tree, use DP on trees to combine longest child paths at each node for diameter.

Common Mistakes

Mistake: Not skipping parent node in DFS causing infinite recursion
Fix: Add check to skip parent node in adjacency traversal

Mistake: Only tracking one longest path instead of two longest paths per node
Fix: Track top two longest child paths to combine for diameter

Summary

Finds the longest path between any two nodes in a tree using dynamic programming.

Use when you need the maximum distance between nodes in a tree structure.

The diameter passes through a node combining its two longest child paths.

Practice

1. What does the diameter of a tree represent in the context of DP on Trees?

easy

A. The longest path between any two nodes in the tree

B. The number of nodes in the tree

C. The height of the tree from the root

D. The sum of all edge weights in the tree

Solution

1. Step 1: Understand the definition of tree diameter

   The diameter is the longest path between any two nodes, not just from the root or sum of edges.

2. Step 2: Differentiate from other tree properties

   Height is from root to leaf, sum of edges is total weight, number of nodes is size, none equal diameter.

3. Final Answer:

   The longest path between any two nodes in the tree -> Option A

4. Quick Check:

   Diameter = Longest path [OK]

Hint: Diameter = longest path between any two nodes [OK]

Common Mistakes:

• Confusing diameter with tree height
• Thinking diameter is total nodes count
• Assuming diameter is sum of all edges

2. Which Rust syntax correctly declares a recursive function to compute the diameter of a tree represented as adjacency list Vec?

easy

A. fn diameter(node: usize, parent: usize, graph: &mut Vec) -> usize { ... }

B. fn diameter(node: usize, graph: Vec) -> usize { ... }

C. fn diameter(node: usize, parent: usize, graph: &Vec) -> usize { ... }

D. fn diameter(&self, node: usize, parent: usize) -> usize { ... }

Solution

1. Step 1: Identify correct parameter passing for graph

   Graph should be passed by reference (&Vec) to avoid copying large data.

2. Step 2: Include parent parameter to avoid revisiting parent node

   Parent node is needed to prevent cycles in undirected tree traversal.

3. Final Answer:

   fn diameter(node: usize, parent: usize, graph: &Vec) -> usize { ... } -> Option C

4. Quick Check:

   Pass graph by reference and track parent [OK]

Hint: Pass graph as &Vec and track parent node [OK]

Common Mistakes:

• Passing graph by value causing performance issues
• Missing parent parameter causing infinite recursion
• Using &mut when mutation is not needed

3. Given the following Rust code snippet to compute the diameter of a tree, what is the printed output?

fn dfs(node: usize, parent: usize, graph: &Vec, max_diameter: &mut usize) -> usize {
    let mut max1 = 0;
    let mut max2 = 0;
    for &child in &graph[node] {
        if child == parent { continue; }
        let height = dfs(child, node, graph, max_diameter);
        if height > max1 {
            max2 = max1;
            max1 = height;
        } else if height > max2 {
            max2 = height;
        }
    }
    *max_diameter = (*max_diameter).max(max1 + max2);
    max1 + 1
}

fn main() {
    let graph = vec![
        vec![1, 2],
        vec![0, 3, 4],
        vec![0],
        vec![1],
        vec![1]
    ];
    let mut max_diameter = 0;
    dfs(0, usize::MAX, &graph, &mut max_diameter);
    println!("{}", max_diameter);
}

medium

A. 3

B. 4

C. 2

D. 5

Solution

1. Step 1: Trace DFS heights from leaves

   Leaves 2, 3, 4 return height 1; node 1 gets max1=1, max2=1; node 0 gets max1=2 (from node 1), max2=1 (from node 2).

2. Step 2: Calculate max diameter updates

   At node 1, max_diameter = max(0,1+1)=2; at node 0, max_diameter = max(2,2+1)=3.

3. Final Answer:

   3 -> Option A

4. Quick Check:

   Longest path length = 3 edges [OK]

Hint: Sum two largest subtree heights at each node [OK]

Common Mistakes:

• Counting nodes instead of edges for diameter
• Not updating max_diameter correctly
• Ignoring parent to avoid revisiting

4. Identify the bug in this Rust function computing tree diameter and select the fix:

fn dfs(node: usize, graph: &Vec, max_diameter: &mut usize) -> usize {
    let mut max1 = 0;
    let mut max2 = 0;
    for &child in &graph[node] {
        let height = dfs(child, graph, max_diameter);
        if height > max1 {
            max2 = max1;
            max1 = height;
        } else if height > max2 {
            max2 = height;
        }
    }
    *max_diameter = (*max_diameter).max(max1 + max2);
    max1 + 1
}

medium

A. Initialize max1 and max2 to 1 instead of 0

B. Add a parent parameter to avoid revisiting the parent node

C. Change max_diameter to return instead of using mutable reference

D. Remove max2 variable as it is unnecessary

Solution

1. Step 1: Identify infinite recursion cause

   Without parent tracking, dfs revisits parent causing infinite recursion or stack overflow.

2. Step 2: Fix by adding parent parameter

   Pass parent node to skip revisiting it in children loop, preventing cycles.

3. Final Answer:

   Add a parent parameter to avoid revisiting the parent node -> Option B

4. Quick Check:

   Parent tracking prevents infinite recursion [OK]

Hint: Always track parent node in undirected tree DFS [OK]

Common Mistakes:

• Not passing parent causing infinite recursion
• Misusing mutable references for max_diameter
• Incorrect initialization of max heights

5. You want to find the diameter of a weighted tree where edges have weights. Which modification to the DP on Trees diameter algorithm is correct?

hard

A. Ignore weights and compute diameter as number of edges only

B. Sum all edge weights and divide by two to get diameter

C. Use BFS twice without any DP or recursion

D. Store and propagate the maximum weighted path length instead of just height counts

Solution

1. Step 1: Understand weighted tree diameter difference

   Diameter must consider edge weights, so heights become max weighted path lengths.

2. Step 2: Modify DP to accumulate weights

   Instead of counting edges, add edge weights when computing subtree heights and update max diameter accordingly.

3. Final Answer:

   Store and propagate the maximum weighted path length instead of just height counts -> Option D

4. Quick Check:

   Weighted diameter needs weighted path sums [OK]

Hint: Accumulate edge weights in DP for weighted diameter [OK]

Common Mistakes:

• Ignoring weights leads to wrong diameter
• Using BFS without weights misses weighted paths
• Incorrect formula dividing total weights