LESSON 305 OF 348 · BINARY TREES

Binary Tree Diameter in Kotlin: One DFS Returning Height and Diameter

Given the root of a binary tree, return the length of the diameter of the tree.

The diameter of a binary tree is the length of the longest path between any two nodes in a tree. This path may or may not pass through the root.

The length of a path between two nodes is represented by the number of edges between them.

Example 1:

Input: root = [1,2,3,4,5]
Output: 3
Explanation: 3 is the length of the path [4,2,1,3] or [5,2,1,3].

Example 2:

Input: root = [1,2]
Output: 1

Constraints:

The number of nodes in the tree is in the range [1, 10⁴].
-100 <= Node.val <= 100

STEP 1 OF 10

Guided

Ctrl + Enter to check

KOTLIN · BINARY TREES

Binary Tree Diameter in Kotlin, explained

Problem overview

The function diameter receives the root of a binary tree, a TreeNode?, and returns the length of its diameter as an Int. The diameter is the longest path between any two nodes. The path may or may not pass through the root, and its length counts edges, not nodes. The tree has between 1 and 10,000 nodes.

In the tree [1, 2, 3, 4, 5], node 1 has children 2 and 3, and node 2 has children 4 and 5. The path 4, 2, 1, 3 has three edges, and so does 5, 2, 1, 3. No path is longer, so the answer is 3. For [1, 2] the only path is from 1 to 2, which is one edge.

Approach: every path bends at one node

Any path in a tree goes up from one node to a highest point and then down to another node. That highest node is where the path "bends". For a bend at a given node, the best path goes down as far as possible on the left and as far as possible on the right, so its length is:

(height of the left side) + (height of the right side)

with heights counted in edges from this node. The diameter of the whole tree is the best value of this over all nodes. A node on the way can also lose to a larger path that lives entirely inside one of its subtrees, which is why the answer for a node is the maximum of three numbers: the diameter of the left subtree, the diameter of the right subtree, and the path that bends here.

The trap is to compute heights separately for every node, which repeats work. Instead, one recursive function returns two values for a subtree: its height and its diameter. A parent combines the pairs from its children in constant time, so each node is visited exactly once.

The example [1, 2, 3, 4, 5] works out like this:

node 4 or 5: leaf              height 0, diameter 0
node 2: children 4 and 5       height 1, diameter 1 + 1 = 2
node 3: leaf                   height 0, diameter 0
node 1: left is 2, right is 3  left side 2, right side 1, diameter max(2, 0, 2 + 1) = 3

Step by step

Return index 1 from diameter. The helper dfs returns an IntArray with the height at index 0 and the diameter at index 1, so the top-level function returns the second entry.
Declare dfs for a nullable TreeNode returning an IntArray of two numbers.
Handle the empty subtree. With no node, both numbers are 0.
Recurse left. Keep the array in left, then copy its height into a var leftHeight and its diameter into a val leftDiameter. The height will change in a moment, and the diameter will not.
Recurse right the same way, with rightHeight and rightDiameter.
Add the edge to the left child. The height a child reports is measured from the child. Reaching it from this node costs one more edge, but only if the child exists, so add 1 when root.left is not null and 0 otherwise.
Do the same on the right.
Compute height as the larger of the two heights: the longest downward path from this node.
Compute diameter as the largest of the left diameter, the right diameter and leftHeight + rightHeight.
Return both in the order the caller expects: height first, diameter second.

The solution in Kotlin

class Solution {
    fun dfs(root: TreeNode?): IntArray {
        if (root == null) {
            return intArrayOf(0, 0)
        }
        val left = dfs(root.left)
        var leftHeight = left[0]
        val leftDiameter = left[1]
        val right = dfs(root.right)
        var rightHeight = right[0]
        val rightDiameter = right[1]
        leftHeight += if (root.left != null) 1 else 0
        rightHeight += if (root.right != null) 1 else 0
        val height = maxOf(leftHeight, rightHeight)
        val diameter = maxOf(maxOf(leftDiameter, rightDiameter), leftHeight + rightHeight)
        return intArrayOf(height, diameter)
    }

    fun diameter(root: TreeNode?): Int {
        return dfs(root)[1]
    }
}

The Kotlin details worth knowing:

Why the conditional + 1. An empty subtree and a single leaf both report height 0, so the parent can't tell them apart from the number alone. That is why the code looks at root.left != null and root.right != null to decide whether an edge exists.
IntArray as a pair. intArrayOf(a, b) builds a primitive int array, read with left[0] and left[1]. It is a compact way to return two numbers. Kotlin also has Pair, and a small data class would name the fields, but the array is enough here. Nothing in the type says which index is which, so the order is a convention.
var versus val. Only the heights change after they are read, so they are var. The diameters and the final results never change and are val.
if is an expression. Kotlin has no ternary operator. if (root.left != null) 1 else 0 produces a value, and += adds it to the variable. The else branch is required when if is used as an expression.
maxOf takes two values here. The call maxOf(a, b) returns the larger of two numbers. For three, as in the diameter line, nest two calls (the standard library has a three-argument overload too, but the nested form is what this solution uses).
Smart cast. After if (root == null) return ..., root is treated as non-null, so root.left compiles without ?..

Complexity

Let n be the number of nodes and h the height of the tree.

Time: O(n). dfs runs once per node and once per missing child. In each call the work is a few comparisons and one small array. The computing of heights is shared between the diameter and the parent, so nothing is recomputed.

Space: O(h). The recursion holds one frame per level of depth, plus a two-element array per call that is soon discarded. A balanced tree has h near log n, and a chain with 10,000 nodes has h = 10,000.

Edge cases and common mistakes

Single node. The diameter is 0, since a path from a node to itself has no edges.
A chain. For [1, 2, null, 3, null, 4], a chain of four nodes each hanging to the left of the one above, the diameter is 3 and the path has no bend.
Counting nodes instead of edges. The path 4, 2, 1, 3 has four nodes but three edges. Counting nodes would give 4.
Assuming the path goes through the root. Take a root with one tiny child on one side and a big subtree on the other. The longest path can lie entirely inside the big subtree, which is why the code also takes the two subtree diameters into account.
Adding 1 for a missing child. Without the null tests, a leaf with no children would appear to have a height of 1 on each side and the result would be too large.
Mixing up the array indices. The height is at index 0 and the diameter at index 1. Reading them the wrong way round gives plausible but wrong numbers.

Testing

[1, 2, 3, 4, 5] gives 3.
[1, 2] gives 1.
[1] gives 0.
[1, 2, null, 3, null, 4] in level order gives 3.

Keep going with KotlinAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free