작심삼일

[LeetCode] 669 | Trim a Binary Search Tree | Python 본문

스터디/코테

[LeetCode] 669 | Trim a Binary Search Tree | Python

yun_s 2022. 4. 15. 09:36
728x90
반응형

문제 링크: https://leetcode.com/problems/trim-a-binary-search-tree/


문제

Given the root of a binary search tree and the lowest and highest boundaries as low and high, trim the tree so that all its elements lies in [low, high]. Trimming the tree should not change the relative structure of the elements that will remain in the tree (i.e., any node's descendant should remain a descendant). It can be proven that there is a unique answer.

Return the root of the trimmed binary search tree. Note that the root may change depending on the given bounds.


조건

  • The number of nodes in the tree in the range [1, $10^4$].
  • 0 <= Node.val <= $10^4$
  • The value of each node in the tree is unique.
  • root is guaranteed to be a valid binary search tree.
  • 0 <= low <= high <= $10^4$

내 풀이

Tree관련 문제는 recursion으로 푸는 것이 도움되는 경우가 종종 있다.

이 문제도 recursion으로 해결하면 쉽다.

BST는 root.left < root < root.right이 항상 성립한다는 것을 염두에 두자.

root.val이 해당 조건 안에 있으면 각각 left와 right을 trimBST로 돌린다.

만약 root.val이 low보다 작으면 root.left는 다 버려야하는 것이므로 root.right만 trimBST를 돌린다.

root.val이 high보다 크면 root.right는 다 버려야하는 것이므로 root.left만 trimBST를 돌린다.


코드

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def trimBST(self, root: Optional[TreeNode], low: int, high: int) -> Optional[TreeNode]:
        if not root:
            return
        
        if low <= root.val <= high:
            root.left = self.trimBST(root.left, low, high)
            root.right = self.trimBST(root.right, low, high)
        elif root.val < low:
            root = self.trimBST(root.right, low, high)
        else:
            root = self.trimBST(root.left, low, high)
        
        return root
728x90
반응형
Comments