leetcode-go/hard/99.go

package main

import (
	"math"
)

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func recoverTree(root *TreeNode) {
	if root == nil {
		return
	}
	var swap1, swap2 *TreeNode
	prev := &TreeNode{math.MinInt32, nil, nil}
	var recurse func(*TreeNode)
	recurse = func(root *TreeNode) {
		if root == nil {
			return
		}
		recurse(root.Left)
		// It's kinda strange, but there's an example:
		// 1 2 6 4 5 3
		// The first element must be the larger one, so pick prev;
		// the second element is the smaller one, pick root.
		if swap1 == nil && root.Val <= prev.Val {
			swap1 = prev
		}
		if swap1 != nil && root.Val <= prev.Val {
			swap2 = root
		}
		prev = root
		recurse(root.Right)
	}
	recurse(root)
	swap1.Val, swap2.Val = swap2.Val, swap1.Val
}

// func main() {
// 	tree := toBinaryTree(nil)
// 	recoverTree(tree)
// 	printTree(tree)
// 	tree = toBinaryTree(newInt(1), newInt(3), nil, nil, newInt(2))
// 	recoverTree(tree)
// 	printTree(tree)
// 	tree = toBinaryTree(newInt(3), newInt(1), newInt(4), nil, nil, newInt(2))
// 	recoverTree(tree)
// 	printTree(tree)
// }