DP on Trees Diameter of Tree in DSA Java

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 between any two towns.

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 to calculate height of each subtree

Heights: 4=0, 5=0 (leaves), 2=?, 3=?, 1=?

Why: We need heights of subtrees to find longest paths through nodes

Step 2: Calculate height of node 2 as max height of children +1

Heights: 4=0, 2=1, 5=0, 3=?, 1=?

Why: Height of node 2 is 1 because child 4 has height 0

Step 3: Calculate height of node 3 as max height of children +1

Heights: 4=0, 2=1, 5=0, 3=1, 1=?

Why: Height of node 3 is 1 because child 5 has height 0

Step 4: Calculate height of node 1 as max height of children +1

Heights: 4=0, 2=1, 5=0, 3=1, 1=2

Why: Height of node 1 is 2 because max child height is 1

Step 5: At each node, calculate diameter candidate as sum of top two child heights + 2 edges

Node 1: top two heights 1 and 1, diameter candidate = 1+1+2=4 edges

Why: Longest path through node 1 is from leaf 4 to leaf 5

Step 6: Compare diameter candidates at all nodes and find maximum

Max diameter = 4 edges = path 4 -> 2 -> 1 -> 3 -> 5

Why: Diameter is longest path found in the tree

Result:

Diameter path: 4 -> 2 -> 1 -> 3 -> 5
Diameter length: 4 edges

Annotated Code

DSA Java

import java.util.*;

public class TreeDiameter {
    static class TreeNode {
        int val;
        List<TreeNode> children = new ArrayList<>();
        TreeNode(int val) { this.val = val; }
    }

    static int diameter = 0;

    // Returns height of subtree rooted at node
    static int dfs(TreeNode node) {
        if (node == null) return -1; // base case for height

        int maxHeight1 = -1; // highest child height
        int maxHeight2 = -1; // second highest child height

        for (TreeNode child : node.children) {
            int h = dfs(child); // get child height

            if (h > maxHeight1) {
                maxHeight2 = maxHeight1;
                maxHeight1 = h;
            } else if (h > maxHeight2) {
                maxHeight2 = h;
            }
        }

        // Update diameter: sum of top two heights + 2 edges
        int candidate = maxHeight1 + maxHeight2 + 2;
        if (candidate > diameter) {
            diameter = candidate;
        }

        return maxHeight1 + 1; // height of current node
    }

    public static void main(String[] args) {
        // Build tree from example: 1-2, 1-3, 2-4, 3-5
        TreeNode n1 = new TreeNode(1);
        TreeNode n2 = new TreeNode(2);
        TreeNode n3 = new TreeNode(3);
        TreeNode n4 = new TreeNode(4);
        TreeNode n5 = new TreeNode(5);

        n1.children.add(n2);
        n1.children.add(n3);
        n2.children.add(n4);
        n3.children.add(n5);

        dfs(n1);
        System.out.println("Diameter length: " + diameter + " edges");
    }
}

for (TreeNode child : node.children) { int h = dfs(child); ... }

recursively find height of each child subtree

if (h > maxHeight1) { maxHeight2 = maxHeight1; maxHeight1 = h; } else if (h > maxHeight2) { maxHeight2 = h; }

track top two maximum child heights

int candidate = maxHeight1 + maxHeight2 + 2; if (candidate > diameter) diameter = candidate;

update diameter with longest path through current node

return maxHeight1 + 1;

return height of current node for parent's calculation

Output

Diameter length: 4 edges

Complexity Analysis

Time: O(n) because we visit each node once in DFS to calculate heights and update diameter

Space: O(n) due to recursion stack in worst case of skewed tree and storing adjacency

vs Alternative: Naive approach checking all pairs of nodes is O(n^2), this DP method is efficient with O(n)

Edge Cases

Single node tree

Diameter is 0 because no edges exist

DSA Java

if (node == null) return -1;

Linear tree (chain)

Diameter equals number of edges in chain, correctly computed by summing heights

DSA Java

int candidate = maxHeight1 + maxHeight2 + 2;

Tree with multiple branches

Diameter is longest path through any node combining two tallest child heights

DSA Java

if (candidate > diameter) diameter = candidate;

When to Use This Pattern

When asked for longest path between any two nodes in a tree, use DP on trees by computing heights and combining top two heights at each node to find diameter.

Common Mistakes

Mistake: Returning height as 0 for null nodes instead of -1, causing off-by-one errors in diameter calculation
Fix: Return -1 for null nodes so leaf height is 0 and edges count correctly

Mistake: Not tracking two highest child heights, only one, missing diameter paths through node
Fix: Track top two maximum child heights to consider paths passing through node

Summary

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

Use when you need the diameter of a tree efficiently without checking all node pairs.

The key insight is the diameter passes through a node combining its two tallest child paths.

Practice

1. What does the diameter of a tree represent in graph theory?

easy

A. The number of leaf nodes in the tree

B. The number of nodes in the tree

C. The height of the tree from the root

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

Solution

1. Step 1: Understand the diameter definition

   The diameter is the longest path between any two nodes in a tree, not just from the root.

2. Step 2: Differentiate from other tree properties

   Height is the longest path from root to leaf, number of nodes and leaves are counts, not paths.

3. Final Answer:

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

4. Quick Check:

   Diameter = Longest path [OK]

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

Common Mistakes:

• Confusing diameter with height
• Thinking diameter counts nodes instead of path length
• Assuming diameter is from root only

2. Which of the following Java code snippets correctly initializes an adjacency list for a tree with n nodes?

easy

A. ArrayList[] tree = new ArrayList[n];

B. List<Integer>[] tree = new ArrayList[n]; for(int i=0; i();

C. List<Integer> tree = new List[n];

D. ArrayList<Integer> tree = new ArrayList<>(n);

Solution

1. Step 1: Understand adjacency list initialization

   We need an array of lists, each initialized to an empty ArrayList to hold neighbors.

2. Step 2: Check each option

   List<Integer>[] tree = new ArrayList[n]; for(int i=0; i(); correctly creates an array of lists and initializes each element. Others either lack initialization or use wrong types.

3. Final Answer:

   List<Integer>[] tree = new ArrayList[n]; for(int i=0; i -> Option B

4. Quick Check:

   Array of lists with initialization [OK]

Hint: Initialize each list in array separately [OK]

Common Mistakes:

• Not initializing each list element
• Using raw List without generics
• Wrong array type declaration

3. Consider the following Java method to find the diameter of a tree. What will be the output if the tree edges are: 0-1, 1-2, 1-3?

int diameter = 0;
int dfs(int node, int parent, List<Integer>[] tree) {
  int max1 = 0, max2 = 0;
  for (int child : tree[node]) {
    if (child == parent) continue;
    int height = dfs(child, node, tree) + 1;
    if (height > max1) {
      max2 = max1;
      max1 = height;
    } else if (height > max2) {
      max2 = height;
    }
  }
  diameter = Math.max(diameter, max1 + max2);
  return max1;
}
// call dfs(0, -1, tree) and then print diameter;

medium

A. 2

B. 3

C. 1

D. 0

Solution

1. Step 1: Build the tree and understand paths

   The tree edges are 0-1, 1-2, 1-3. The longest path is from node 2 to node 0 or 3 to 0, length 2 edges.

2. Step 2: Trace DFS and diameter calculation

   DFS calculates max heights from children and updates diameter as sum of two largest heights. Here, diameter becomes 2.

3. Final Answer:

   2 -> Option A

4. Quick Check:

   Longest path length = 2 edges [OK]

Hint: Diameter = sum of two largest subtree heights [OK]

Common Mistakes:

• Counting nodes instead of edges
• Ignoring parent check in DFS
• Miscomputing max heights

4. The following code attempts to compute the diameter of a tree but has a bug. What is the bug?

int diameter = 0;
int dfs(int node, int parent, List<Integer>[] tree) {
  int max1 = 0, max2 = 0;
  for (int child : tree[node]) {
    int height = dfs(child, node, tree) + 1;
    if (height > max1) {
      max2 = max1;
      max1 = height;
    } else if (height > max2) {
      max2 = height;
    }
  }
  diameter = Math.max(diameter, max1 + max2);
  return max1;
}
// call dfs(0, -1, tree);

medium

A. Missing check to skip parent node in DFS

B. Incorrect update of diameter variable

C. Wrong return value from DFS

D. Initialization of max1 and max2 is incorrect

Solution

1. Step 1: Analyze DFS loop over children

   The code calls dfs on all children but does not skip the parent node, causing infinite recursion or stack overflow.

2. Step 2: Identify missing condition

   There should be a condition to continue if child == parent to avoid revisiting the parent.

3. Final Answer:

   Missing check to skip parent node in DFS -> Option A

4. Quick Check:

   Skip parent to prevent infinite recursion [OK]

Hint: Always skip parent node in tree DFS [OK]

Common Mistakes:

• Forgetting to skip parent node
• Misplacing diameter update
• Returning wrong height

5. You want to find the diameter of a weighted tree where edges have positive weights. Which modification to the standard diameter DFS approach is correct?

hard

A. Use BFS twice without considering weights

B. Ignore weights and count edges as 1 each

C. Add edge weights to the height calculation instead of just adding 1

D. Calculate diameter by summing all edge weights in the tree

Solution

1. Step 1: Understand weighted tree diameter

   In weighted trees, path length is sum of edge weights, not just number of edges.

2. Step 2: Modify DFS height calculation

   Instead of adding 1 for each edge, add the actual edge weight to the height returned by DFS.

3. Step 3: Update diameter similarly

   Diameter is max sum of two largest weighted heights from children.

4. Final Answer:

   Add edge weights to the height calculation instead of just adding 1 -> Option C

5. Quick Check:

   Use weights in height sums for weighted diameter [OK]

Hint: Add edge weights, not just count edges [OK]

Common Mistakes:

• Ignoring weights in diameter calculation
• Using BFS without weights
• Summing all edges instead of longest path