Diagonal Traversal of a Binary tree

The problem, simply

You get a binary tree. Draw slanted lines through it, going down and to the right, like \. Each line is a diagonal.

• The root is on diagonal 0.
• A right child is on the same diagonal as its parent.
• A left child is on the next diagonal, one more than its parent.

List the nodes diagonal by diagonal: all of diagonal 0 first, then diagonal 1, and so on. Inside a diagonal, list nodes from top to bottom, the way they are met when you read the tree.

A story to remember it

Picture it

Examples

Way 1: the simple way

Walk the tree with a deep dive, carrying the diagonal number. Add each node to a list for its diagonal in a sorted map. Visit the node first, then the left side, then the right side, so that nodes higher up land in their list first.

1. Start at the root with diagonal 0.
2. Add the node to the list for its diagonal.
3. Visit the left child with diagonal + 1, then the right child with the same diagonal.
4. Read the lists from the smallest diagonal to the biggest.

import java.util.*;

class Solution {
    public List<Integer> diagonal(TreeNode root) {
        TreeMap<Integer, List<Integer>> lines = new TreeMap<>();
        dive(root, 0, lines);
        List<Integer> out = new ArrayList<>();
        for (List<Integer> line : lines.values()) out.addAll(line);
        return out;
    }

    private void dive(TreeNode n, int d, TreeMap<Integer, List<Integer>> lines) {
        if (n == null) return;
        lines.computeIfAbsent(d, k -> new ArrayList<>()).add(n.val);
        dive(n.left, d + 1, lines);    // left first keeps the expected order
        dive(n.right, d, lines);
    }
}

How long it takes: each node is visited once, and each visit talks to a sorted map, O(n log n). (An ArrayList of lists indexed by diagonal would make it O(n).) Extra room: the map and the dive, O(n).

Way 2: the clever way

Use the water park line, a queue. Each ticket in the line starts a slide.

1. Put the root in the line.
2. Take a node from the front. Slide right from it: write down every node on the way.
3. While sliding, every node with a left child gives that child a ticket at the back of the line.
4. When the slide ends, take the next ticket. Repeat until the line is empty.

Tickets for diagonal 1 are all handed out while sliding diagonal 0, and so on. So the line serves the diagonals in order without counting them.

import java.util.*;

class Solution {
    public List<Integer> diagonal(TreeNode root) {
        List<Integer> out = new ArrayList<>();
        if (root == null) return out;
        Queue<TreeNode> line = new ArrayDeque<>();
        line.add(root);
        while (!line.isEmpty()) {
            TreeNode n = line.poll();          // start a new slide
            while (n != null) {
                out.add(n.val);
                if (n.left != null) line.add(n.left);   // a ticket for later
                n = n.right;                   // keep sliding right
            }
        }
        return out;
    }
}

How long it takes: every node is written once and joins the line at most once, O(n). Extra room: the line, O(n) in the worst case.

Watch out for

• The direction of the slant. Here diagonals go down to the right, so right children stay on the same line. Some problems slant the other way, where left children stay.
• Wanting the diagonals as separate lists? In the queue version, take size = line.size() before each round and run that many slides per round. Each round is one diagonal.
• An empty tree. Return an empty list.

What interviewers ask next

• Sum of each diagonal? Same walk, but add values into a running total per diagonal.
• Why does the queue keep the diagonals in order? Every ticket for diagonal d + 1 is handed out while sliding diagonal d, and all of diagonal d was in line before. So the line never mixes them up.
• How does this relate to vertical order? Vertical order uses columns (left is minus 1, right is plus 1). Diagonal order uses "left is plus 1, right is plus 0". The tool is the same: give each node a number and group by it.
• Anti diagonals (slanting the other way)? Swap the roles: slide left and give tickets to right children.

Remember this

Grown-up words: the slides are diagonals, the waiting line is a queue, the tickets are left children added to the queue, and the slanted grouping is the diagonal traversal.