LeetCode/go/diameter-of-binary-tree.go

package main

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

func diameterAndHeight(node *TreeNode) (int, int) {
	if node == nil {
		return 0, 0
	}

	leftDiameter, leftHeight := diameterAndHeight(node.Left)
	rightDiameter, rightHeight := diameterAndHeight(node.Right)

	diameter := max(leftHeight+rightHeight, max(leftDiameter, rightDiameter))
	height := 1 + max(leftHeight, rightHeight)

	return diameter, height
}

func diameterOfBinaryTree(root *TreeNode) int {
	diameter, _ := diameterAndHeight(root)
	return diameter
}