[LeetCode] 700 | Search in an Binary Search Tree | Python

---

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

---

문제

You are given the root of a binary search tree (BST) and an integer val.

Find the node in the BST that the node's value equals val and return the subtree rooted with that node. If such a node does not exist, return null.

---

조건

• The number of nodes in the tree is in the range [1, 5000].
• 1 <= Node.val <= $10^7$
• root is a binary search tree.
• 1 <= val <= $10^7$

---

내 풀이

BST의 특성을 이용해서 root가 None이 아닐 동안 while문이 계속 돌도록 한다.

while문 안에서는 val이 root.val보다 작을 경우 왼쪽으로 가야하므로 root.left로 보내고, root.val보다 클 경우 오른쪽(root.right)으로 보낸다.

다만 위 경우에서 root.left나 root.right이 존재하지 않으면 null을 리턴하면 된다.

val이 root.val과 일치하면 현재 root를 리턴하면 된다.

---

코드

# 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 searchBST(self, root: Optional[TreeNode], val: int) -> Optional[TreeNode]:
        while root:
            if root.val < val:
                if root.right:
                    root = root.right
                else:
                    return
            elif root.val > val:
                if root.left:
                    root = root.left
                else:
                    return
            else:
                return root