DP on Trees Diameter of Tree in DSA Go

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 cities (nodes). The diameter is the longest route you can drive without repeating roads. We find it by looking at the longest roads from each city and combining them.

    1
   / \
  2   3
     / \
    4   5

Nodes connected like a tree with branches.

Dry Run Walkthrough

Input: Tree edges: 1-2, 1-3, 3-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 (leaf), longest path down is 0

Node 2: longest_down = 0

Why: Leaf nodes have no children, so longest path down is zero

Step 3: At node 3, explore children 4 and 5

Node 3: children 4 and 5 to explore

Why: Calculate longest paths down from children to find node 3's longest_down

Step 4: At node 4 (leaf), longest path down is 0

Node 4: longest_down = 0

Why: Leaf node

Step 5: At node 5 (leaf), longest path down is 0

Node 5: longest_down = 0

Why: Leaf node

Step 6: At node 3, longest_down = max(0,0) + 1 = 1; diameter through 3 = 0 + 0 + 2 = 2

Node 3: longest_down=1, diameter_candidate=2

Why: Longest path down is one edge plus max child's longest_down; diameter through node is sum of top two longest_down + 2 edges

Step 7: At node 1, longest_down = max(0,1) + 1 = 2; diameter through 1 = 0 + 1 + 2 = 3

Node 1: longest_down=2, diameter_candidate=3

Why: Combine longest_down from children 2 and 3 plus edges to get diameter candidate

Result:

Diameter is 3 edges: path 2 -> 1 -> 3 -> 4 (or 5)

Annotated Code

DSA Go

package main

import "fmt"

// TreeNode represents a node in the tree
type TreeNode struct {
	val      int
	children []*TreeNode
}

var diameter int

// dfs returns the longest path down from current node
func dfs(node *TreeNode) int {
	if node == nil {
		return 0
	}
	max1, max2 := 0, 0
	// check all children
	for _, child := range node.children {
		length := dfs(child) + 1
		// update top two longest lengths
		if length > max1 {
			max2 = max1
			max1 = length
		} else if length > max2 {
			max2 = length
		}
	}
	// update global diameter if sum of two longest paths is bigger
	if max1+max2 > diameter {
		diameter = max1 + max2
	}
	return max1
}

func main() {
	// Build tree from example edges: 1-2, 1-3, 3-4, 3-5
	node1 := &TreeNode{val: 1}
	node2 := &TreeNode{val: 2}
	node3 := &TreeNode{val: 3}
	node4 := &TreeNode{val: 4}
	node5 := &TreeNode{val: 5}

	node1.children = []*TreeNode{node2, node3}
	node3.children = []*TreeNode{node4, node5}

	diameter = 0
	dfs(node1)
	fmt.Printf("Diameter of tree: %d\n", diameter)
}

for _, child := range node.children {

iterate over all children to find longest paths down

length := dfs(child) + 1

calculate longest path down from child plus edge to child

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

keep track of top two longest child paths

if max1+max2 > diameter { diameter = max1 + max2 }

update global diameter if current node's two longest paths sum is bigger

return max1

return longest path down from current node

Output

Diameter of tree: 3

Complexity Analysis

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

Space: O(n) due to recursion stack in worst case (skewed tree)

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

Edge Cases

Empty tree (no nodes)

Diameter is 0 because no paths exist

DSA Go

if node == nil { return 0 }

Single node tree

Diameter is 0 because no edges exist

DSA Go

if node == nil { return 0 }

Linear tree (chain)

Diameter equals number of edges in chain

DSA Go

updates diameter with max1+max2

When to Use This Pattern

When asked for longest path between any two nodes in a tree, use DP on trees to find longest downward paths and combine them for diameter.

Common Mistakes

Mistake: Only tracking one longest path from children, missing second longest for diameter
Fix: Track top two longest child paths to compute diameter through current node

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 downward paths.

Practice

1. What is the diameter of a tree in the context of dynamic programming 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 number of leaf nodes in the tree

Solution

1. Step 1: Understand the definition of tree diameter

   The diameter is the longest path between any two nodes in the tree, not necessarily passing through the root.

2. Step 2: Differentiate from other tree properties

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

3. Final Answer:

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

4. Quick Check:

   Diameter = longest path in tree [OK]

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

Common Mistakes:

• Confusing diameter with height
• Thinking diameter is number of nodes
• Assuming diameter passes through root

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

easy

A. graph := make([][]int, n)

B. graph := make([]int, n)

C. graph := make(map[int]int)

D. graph := make([]int, 0)

Solution

1. Step 1: Identify adjacency list structure

   An adjacency list for a tree with n nodes is a slice of slices of ints, each inner slice holds neighbors.

2. Step 2: Match Go syntax for adjacency list

   make([][]int, n) creates a slice of n slices of ints, perfect for adjacency list initialization.

3. Final Answer:

   graph := make([][]int, n) -> Option A

4. Quick Check:

   Adjacency list = [][]int slice [OK]

Hint: Adjacency list = slice of slices in Go [OK]

Common Mistakes:

• Using make([]int, n) which is 1D slice
• Using map[int]int which is not adjacency list
• Initializing empty slice without size

3. Given the following Go code to find the diameter of a tree, what is the output?

package main
import "fmt"

var graph [][]int
var maxDiameter int

func dfs(node, parent int) int {
    max1, max2 := 0, 0
    for _, child := range graph[node] {
        if child == parent {
            continue
        }
        length := dfs(child, node) + 1
        if length > max1 {
            max2 = max1
            max1 = length
        } else if length > max2 {
            max2 = length
        }
    }
    if max1+max2 > maxDiameter {
        maxDiameter = max1 + max2
    }
    return max1
}

func main() {
    graph = make([][]int, 5)
    graph[0] = []int{1, 2}
    graph[1] = []int{0, 3}
    graph[2] = []int{0}
    graph[3] = []int{1, 4}
    graph[4] = []int{3}
    dfs(0, -1)
    fmt.Println(maxDiameter)
}

medium

A. 2

B. 4

C. 5

D. 3

Solution

1. Step 1: Understand the tree structure

   The tree edges are: 0-1, 0-2, 1-3, 3-4. The longest path is from node 2 to node 4 via 0,1,3 with length 4 edges.

2. Step 2: Trace DFS and diameter calculation

   DFS calculates max1 and max2 longest paths from children. maxDiameter updates to max1+max2 which is 4 for path 2-0-1-3-4.

3. Final Answer:

   4 -> Option B

4. Quick Check:

   Longest path length = 4 edges [OK]

Hint: Diameter = longest path edges counted in DFS [OK]

Common Mistakes:

• Counting nodes instead of edges
• Ignoring parent check in DFS
• Misreading tree connections

4. The following Go code attempts to compute the diameter of a tree but produces incorrect results. What is the bug?

func dfs(node, parent int) int {
    max1, max2 := 0, 0
    for _, child := range graph[node] {
        length := dfs(child, node) + 1
        if length > max1 {
            max2 = max1
            max1 = length
        } else if length > max2 {
            max2 = length
        }
    }
    if max1+max2 > maxDiameter {
        maxDiameter = max1 + max2
    }
    return max1
}

medium

A. Wrong update of maxDiameter variable

B. Incorrect initialization of max1 and max2

C. Missing check to skip parent node in DFS

D. Returning max2 instead of max1

Solution

1. Step 1: Analyze DFS loop over children

   The code calls dfs on all children without skipping the parent, causing infinite recursion or wrong paths.

2. Step 2: Identify missing parent check

   To avoid revisiting parent, code must skip child == parent in loop.

3. Final Answer:

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

4. Quick Check:

   Skip parent in DFS to avoid cycles [OK]

Hint: Always skip parent node in DFS to prevent cycles [OK]

Common Mistakes:

• Forgetting parent skip causes infinite recursion
• Confusing max1 and max2 updates
• Incorrect diameter update logic

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

hard

A. Use BFS twice without considering weights

B. Ignore edge weights and count edges as 1 each

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

D. Replace all increments of 1 with the edge weight when calculating path lengths

Solution

1. Step 1: Understand weighted tree diameter

   Diameter is longest weighted path, so edge weights must be included in path length calculations.

2. Step 2: Modify DP algorithm for weights

   Instead of adding 1 for each edge, add the actual edge weight when updating max1 and max2.

3. Step 3: Reject incorrect options

   Ignoring weights or summing all edges does not find longest path. BFS without weights finds unweighted diameter only.

4. Final Answer:

   Replace all increments of 1 with the edge weight when calculating path lengths -> Option D

5. Quick Check:

   Use edge weights in path sums for weighted diameter [OK]

Hint: Add edge weights instead of 1 for weighted diameter [OK]

Common Mistakes:

• Ignoring weights in diameter calculation
• Summing all edges instead of longest path
• Using unweighted BFS for weighted tree