Update trim-a-binary-search-tree.py

kamyu104 · web-flow · commit 221c94db63b6 · 2017-09-04T12:57:09.000+08:00
diff --git a/Python/trim-a-binary-search-tree.py b/Python/trim-a-binary-search-tree.py
@@ -1,29 +1,65 @@
-// Time:  O(n)
-// Space: O(h)
+# Time:  O(n)
+# Space: O(h)
 
-/**
- * Definition for a binary tree node.
- * struct TreeNode {
- *     int val;
- *     TreeNode *left;
- *     TreeNode *right;
- *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
- * };
- */
-class Solution {
-public:
-    TreeNode* trimBST(TreeNode* root, int L, int R) {
-        if (!root) {
-            return nullptr;
-        }
-        if (root->val < L) {
-            return trimBST(root->right, L, R);
-        }
-        if (root->val > R) {
-            return trimBST(root->left, L, R);
-        }
-        root->left = trimBST(root->left, L, R);
-        root->right = trimBST(root->right, L, R);
-        return root;
-    }
-};
+# Given a binary search tree and the lowest and highest boundaries as L and R,
+# trim the tree so that all its elements lies in [L, R] (R >= L).
+# You might need to change the root of the tree, so the result should
+# return the new root of the trimmed binary search tree.
+#
+# Example 1:
+# Input: 
+#     1
+#    / \
+#   0   2
+#
+#   L = 1
+#   R = 2
+#
+# Output: 
+#     1
+#       \
+#        2
+# Example 2:
+# Input: 
+#     3
+#    / \
+#   0   4
+#    \
+#     2
+#    /
+#   1
+#
+#   L = 1
+#   R = 3
+#
+# Output: 
+#       3
+#      / 
+#    2   
+#   /
+#  1
+
+# 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 trimBST(self, root, L, R):
+        """
+        :type root: TreeNode
+        :type L: int
+        :type R: int
+        :rtype: TreeNode
+        """
+        if not root:
+            return None
+        if root.val < L:
+            return self.trimBST(root.right, L, R)
+        if root.val > R:
+            return self.trimBST(root.left, L, R)
+        root.left, root.right = self.trimBST(root.left, L, R), self.trimBST(root.right, L, R)
+        return root
+    
