LWC 70: 776. Split BST

LWC 70: 776. Split BST

傳送門：776. Split BST

Problem:

Given a Binary Search Tree (BST) with root node root, and a target value V, split the tree into two subtrees where one subtree has nodes that are all smaller or equal to the target value, while the other subtree has all nodes that are greater than the target value. It’s not necessarily the case that the tree contains a node with value V.

Additionally, most of the structure of the original tree should remain. Formally, for any child C with parent P in the original tree, if they are both in the same subtree after the split, then node C should still have the parent P.

You should output the root TreeNode of both subtrees after splitting, in any order.

Example 1:

Input: root = [4,2,6,1,3,5,7], V = 2
Output: [[2,1],[4,3,6,null,null,5,7]]
Explanation:
Note that root, output[0], and output1 are TreeNode objects, not arrays.

The given tree [4,2,6,1,3,5,7] is represented by the following diagram:

      4
    /   \
  2      6
 / \    / \
1   3  5   7

while the diagrams for the outputs are:

      4
    /   \
  3      6      and    2
        / \           /
       5   7         1

Note:

• The size of the BST will not exceed 50.
• The BST is always valid and each node’s value is different.

思路：
問題的關鍵在於進行切分後，樹的結構不能改變。影響BST的結構在於輸入順序，所以切分前後，維持輸入順序即可。而BST的層序遍歷剛好是最初的輸入順序。所以
1. 求出輸入順序。
2. 根據val劃分兩個子輸入順序
3. 根據子輸入順序建立BST

代碼如下：

    public TreeNode[] splitBST(TreeNode root, int V) {
        if (root == null) return new TreeNode[] {null, null};
        List<Integer> nodes = new ArrayList<>();
        Queue<TreeNode> q = new ArrayDeque<>();
        q.offer(root);
        while (!q.isEmpty()) {
            TreeNode now = q.poll();
            nodes.add(now.val);
            if (now.left != null) {
                q.offer(now.left);
            }
            if (now.right != null) {
                q.offer(now.right);
            }
        }
        List<Integer> left = new ArrayList<>();
        List<Integer> right = new ArrayList<>();
        for (int val : nodes) {
            if (val <= V) left.add(val);
            else right.add(val);
        }

        TreeNode leftNode = null;
        for (int val : left) leftNode = build(leftNode, val);

        TreeNode rightNode = null;
        for (int val : right) rightNode = build(rightNode, val);

        return new TreeNode[] {leftNode, rightNode};
    }

    TreeNode build(TreeNode root, int val) {
        if (root == null) return new TreeNode(val);
        else {
            if (val <= root.val) {
                root.left = build(root.left, val);
            }
            else {
                root.right = build(root.right, val);
            }
            return root;
        }
    }

再來一種遞歸的思路，很簡單。我們把問題進行拆分合並式求解。比如當root.val <= val時，說明root和root的左子樹均小於等於val，這種結構維持，且以root爲根，那麼問題就轉而求root右子樹的分裂。因爲分裂的第一顆樹，是小於val的，所以需要鏈接到root的右子樹上，詳見代碼。

Java版本：

    public TreeNode[] splitBST(TreeNode root, int V) {
        if (root == null) return new TreeNode[] {null, null};
        if (root.val <= V) {
            TreeNode[] res = splitBST(root.right, V);
            root.right = res[0];
            res[0] = root;
            return res;
        }
        else{
            TreeNode[] res = splitBST(root.left, V);
            root.left = res[1];
            res[1] = root;
            return res;
        }
    }

Python版本：

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def splitBST(self, root, V):
        """
        :type root: TreeNode
        :type V: int
        :rtype: List[TreeNode]
        """
        if not root: return [None] * 2
        if root.val <= V:
            res = self.splitBST(root.right, V)
            root.right = res[0]
            res[0] = root
            return res
        else:
            res = self.splitBST(root.left, V)
            root.left = res[1]
            res[1] = root
            return res