leetcode — flatten-binary-tree-to-linked-list

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * Source : https://oj.leetcode.com/problems/flatten-binary-tree-to-linked-list/
 *
 *
 * Given a binary tree, flatten it to a linked list in-place.
 *
 * For example,
 * Given
 *
 *          1
 *         / \
 *        2   5
 *       / \   \
 *      3   4   6
 *
 * The flattened tree should look like:
 *
 *    1
 *     \
 *      2
 *       \
 *        3
 *         \
 *          4
 *           \
 *            5
 *             \
 *              6
 *
 *
 * Hints:
 * If you notice carefully in the flattened tree, each node's right child points to
 * the next node of a pre-order traversal.
 */
public class FlattenBinaryTree {

    /**
     * 把一棵二叉树按照preoorder展开为一个单向链表
     *
     * 先使用preorder traversal将树遍历到一个list中，然后按顺序连接起来
     *
     * @param root
     * @return
     */
    public List<TreeNode> flatten (TreeNode root) {
        if (root == null) {
            return null;
        }
        List<TreeNode> list = new ArrayList<TreeNode>();
        flattenByRecursion(root, list);
        for (int i = 0; i < list.size() - 1; i++) {
            list.get(i).leftChild = null;
            list.get(i).rightChild = list.get(i+1);
        }
        list.get(list.size()-1).leftChild = null;
        list.get(list.size()-1).rightChild = null;
        return list;
    }

    public void flattenByRecursion (TreeNode root, List<TreeNode> list) {
        if (root == null) {
            return;
        }
        list.add(root);
        flattenByRecursion(root.leftChild, list);
        flattenByRecursion(root.rightChild, list);
    }

    /**
     * 上面的方法好在简单，但是空间复杂度是O(n)，还可以优化以下占用的空间
     *
     * 使用常数空间来完成flattern
           1
          / \
         2   5
         \    \
         3    6 <- rightTail
         \
         4  <- leftTail

     * 假设root的左右子树都已经flatten，那么需要操作就是
     * temp = root.right
     * root.right = root.left
     * leftTail.right = temp
     * root.left = null
     *
     * 递归使用这种方法
     * 这种情况下递归的时候需要返回tail
     *
     * 递归的时候需要返回tail节点
     *
     * @param root
     * @return
     */
    public TreeNode flattenByRecursion1 (TreeNode root) {
        if (root == null) {
            return null;
        }
        TreeNode leftTail = flattenByRecursion1(root.leftChild);
        TreeNode rightTail = flattenByRecursion1(root.rightChild);
        if (root.leftChild != null) {
            TreeNode temp = root.rightChild;
            root.rightChild = root.leftChild;
            root.leftChild = null;
            leftTail.rightChild = temp;
        }
        // 因为上面已经把左子树拼接在了root和root.right之间，所以先判断顺序是：是否有右子树，是否有左子树
        // 右子树不为空：如果右子树不为空，表示右子树才是尾节点，因为左子树的尾节点还在右子树的上面
        if (rightTail != null) {
            return rightTail;
        }
        // 右子树为空，左子树不为空
        if (leftTail != null) {
            return leftTail;
        }
        // 左右子树都为空：返回root
        return root;
    }

    public TreeNode flatten1(TreeNode root) {
        flattenByRecursion1(root);
        return root;
    }

    public TreeNode createTree (char[] treeArr) {
        TreeNode[] tree = new TreeNode[treeArr.length];
        for (int i = 0; i < treeArr.length; i++) {
            if (treeArr[i] == '#') {
                tree[i] = null;
                continue;
            }
            tree[i] = new TreeNode(treeArr[i]-'0');
        }
        int pos = 0;
        for (int i = 0; i < treeArr.length && pos < treeArr.length-1; i++) {
            if (tree[i] != null) {
                tree[i].leftChild = tree[++pos];
                if (pos < treeArr.length-1) {
                    tree[i].rightChild = tree[++pos];
                }
            }
        }
        return tree[0];
    }

    public static void print (TreeNode root) {
        List<TreeNode> list = new ArrayList<TreeNode>();
        while (root != null) {
            list.add(root);
            root = root.rightChild;
        }
        System.out.println(Arrays.toString(list.toArray(new TreeNode[list.size()])));
    }

    private class TreeNode {
        TreeNode leftChild;
        TreeNode rightChild;
        int value;

        public TreeNode(int value) {
            this.value = value;
        }

        public TreeNode() {

        }

        @Override
        public String toString() {
            return value + "";
        }
    }

    public static void main(String[] args) {
        FlattenBinaryTree flattenBinaryTree = new FlattenBinaryTree();
        char[] arr = new char[]{'1','2','5','3','4','#','6'};

        List<TreeNode> list = flattenBinaryTree.flatten(flattenBinaryTree.createTree(arr));

        System.out.println(Arrays.toString(list.toArray(new TreeNode[list.size()])));

        print(flattenBinaryTree.flatten1(flattenBinaryTree.createTree(arr)));

    }
}