LeetCode in Kotlin

114. Flatten Binary Tree to Linked List

Medium

Given the root of a binary tree, flatten the tree into a “linked list”:

Example 1:

Input: root = [1,2,5,3,4,null,6]

Output: [1,null,2,null,3,null,4,null,5,null,6]

Example 2:

Input: root = []

Output: []

Example 3:

Input: root = [0]

Output: [0]

Constraints:

Follow up: Can you flatten the tree in-place (with O(1) extra space)?

Solution

import com_github_leetcode.TreeNode

/*
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    fun flatten(root: TreeNode?) {
        if (root != null) {
            findTail(root)
        }
    }

    private fun findTail(root: TreeNode): TreeNode {
        val left: TreeNode? = root.left
        val right: TreeNode? = root.right
        val tail: TreeNode
        // find the tail of left subtree, tail means the most left leaf
        if (left != null) {
            tail = findTail(left)
            // stitch the right subtree below the tail
            root.left = null
            root.right = left
            tail.right = right
        } else {
            tail = root
        }
        // find tail of the right subtree
        return if (tail.right == null) {
            tail
        } else {
            findTail(tail.right!!)
        }
    }
}