987. Vertical Order Traversal of a Binary Tree #

Problem #

Given the root of a binary tree, calculate the vertical order traversal of the binary tree.

For each node at position (row, col), its left and right children will be at positions (row + 1, col - 1) and (row + 1, col + 1) respectively. The root of the tree is at (0, 0).

The vertical order traversal of a binary tree is a list of top-to-bottom orderings for each column index starting from the leftmost column and ending on the rightmost column. There may be multiple nodes in the same row and same column. In such a case, sort these nodes by their values.

Return the vertical order traversal of the binary tree.

Example 1:

Input: root = [3,9,20,null,null,15,7]
Output: [[9],[3,15],[20],[7]]
Explanation:
Column -1: Only node 9 is in this column.
Column 0: Nodes 3 and 15 are in this column in that order from top to bottom.
Column 1: Only node 20 is in this column.
Column 2: Only node 7 is in this column.

Example 2:

Input: root = [1,2,3,4,5,6,7]
Output: [[4],[2],[1,5,6],[3],[7]]
Explanation:
Column -2: Only node 4 is in this column.
Column -1: Only node 2 is in this column.
Column 0: Nodes 1, 5, and 6 are in this column.
          1 is at the top, so it comes first.
          5 and 6 are at the same position (2, 0), so we order them by their value, 5 before 6.
Column 1: Only node 3 is in this column.
Column 2: Only node 7 is in this column.

Example 3:

Input: root = [1,2,3,4,6,5,7]
Output: [[4],[2],[1,5,6],[3],[7]]
Explanation:
This case is the exact same as example 2, but with nodes 5 and 6 swapped.
Note that the solution remains the same since 5 and 6 are in the same location and should be ordered by their values.

Constraints:

The number of nodes in the tree is in the range [1, 1000].
0 <= Node.val <= 1000

Problem Summary #

Given the root node root of a binary tree, design an algorithm to calculate the vertical order traversal sequence of the binary tree.

For each node located at (row, col), its left and right child nodes are located at (row + 1, col - 1) and (row + 1, col + 1), respectively. The root node of the tree is located at (0, 0). The vertical order traversal of a binary tree starts from the leftmost column and ends at the rightmost column, collecting all nodes in each column by column index to form an ordered list sorted from top to bottom according to their positions. If there are multiple nodes in the same row and column, sort them by node value in ascending order. Return the vertical order traversal sequence of the binary tree.

Solution Approach #

The problem requires traversing the binary tree column by column. Two issues need to be solved. The first is how to calculate the two-dimensional coordinates of each node in the binary tree. The second is that when multiple nodes are stacked at the same two-dimensional coordinate point, they need to be sorted in ascending order, as in Example 2 and Example 3, where at the same two-dimensional coordinate point (2, 0), there are 2 different nodes stacked together.
First solve the first issue. Since the problem states that the root node is (0, 0), meaning the root node is the coordinate origin, the x-coordinates of its left subtree are all negative, and the x-coordinates of its right subtree are all positive. Using preorder traversal, we can calculate the two-dimensional coordinates of these nodes. Then perform a sort: sort by x-coordinate in ascending order; when coordinates are the same, this corresponds to the case where nodes are stacked together, and the stacked nodes are sorted by val in ascending order. In this way, all nodes are arranged along the x-axis. After sorting is complete, the second issue is also solved.
The final step only needs to scan this sorted array once, and in column order, package the nodes in the same column into a one-dimensional array one by one. The resulting two-dimensional array is the answer required by the problem.

Code #

package leetcode

import (
	"math"
	"sort"

	"github.com/halfrost/leetcode-go/structures"
)

// TreeNode define
type TreeNode = structures.TreeNode

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

type node struct {
	x, y, val int
}

func verticalTraversal(root *TreeNode) [][]int {
	var dfs func(root *TreeNode, x, y int)
	var nodes []node
	dfs = func(root *TreeNode, x, y int) {
		if root == nil {
			return
		}
		nodes = append(nodes, node{x, y, root.Val})
		dfs(root.Left, x+1, y-1)
		dfs(root.Right, x+1, y+1)
	}
	dfs(root, 0, 0)

	sort.Slice(nodes, func(i, j int) bool {
		a, b := nodes[i], nodes[j]
		return a.y < b.y || a.y == b.y &&
			(a.x < b.x || a.x == b.x && a.val < b.val)
	})

	var res [][]int
	lastY := math.MinInt32
	for _, node := range nodes {
		if lastY != node.y {
			res = append(res, []int{node.val})
			lastY = node.y
		} else {
			res[len(res)-1] = append(res[len(res)-1], node.val)
		}
	}
	return res
}