1261. Find Elements in a Contaminated Binary Tree

Find Elements in a Contaminated Binary Tree - LeetCode

-1로 값이 채워진 이진 탐색 트리가 주어진다.

주어진 조건에 맞게 이진 탐색 트리의 값을 복원한 뒤 반환하려 할 때, 아래 코드를 완성시키는 것이 요구사항이다.

 

조건

1. root.val == 0
2. If treeNode.val == x and treeNode.left != null, then treeNode.left.val == 2 * x + 1
3. If treeNode.val == x and treeNode.right != null, then treeNode.right.val == 2 * x + 2

완성해야 하는 코드

# 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 FindElements:

    def __init__(self, root: Optional[TreeNode]):
    	# 자유롭게 사용

    def find(self, target: int) -> bool:
    	# 복원한 트리에서 target 값을 가지는 노드가 존재하면 True, 존재하지 않는다면 False 반환
        


# Your FindElements object will be instantiated and called as such:
# obj = FindElements(root)
# param_1 = obj.find(target)

예시

왼쪽 트리가 주어지면 오른쪽 트리를 반환하면 된다.

1. DFS

from typing import *


class FindElements:

    def __init__(self, root: Optional[TreeNode]):
        self.root = root
        self.root.val = 0
        self.visited = {0}
        self.recover(self.root)

    def find(self, target: int) -> bool:
        if target in self.visited:
            return True
        return False

    def recover(self, cur):
        if cur.left:
            cur.left.val = cur.val * 2 + 1
            self.visited.add(cur.left.val)
            self.recover(cur.left)
        if cur.right:
            cur.right.val = cur.val * 2 + 2
            self.visited.add(cur.right.val)
            self.recover(cur.right)

우선 트리는 만들어서 주어지니 트리 아래로 탐색을 진행하면서 조건대로 값을 변경하고, 변경한 값을 set에 저장한다.

find()는 트리에 있는 모든 노드의 값을 미리 저장했기 때문에 O(1)로 바로바로 찾을 수 있다.