1305. All Elements in Two Binary Search Trees

All Elements in Two Binary Search Trees - LeetCode

2개의 이진 탐색 트리가 주어진다.

두 트리의 value값을 오름차순 정렬한 list를 반환하는 문제다.

 

예시

Input: root1 = [2,1,4], root2 = [1,0,3]
Output: [0,1,1,2,3,4]

 

1. 중위 순회(in-order traversal)

class Solution:
    def getAllElements(self, root1: TreeNode, root2: TreeNode) -> List[int]:
        root1_seq = deque([])
        root2_seq = deque([])
        answer = []

        def in_order(cur, flag):
            if cur.left:
                in_order(cur.left, flag)
            if flag == 1:
                root1_seq.append(cur.val)
            else:
                root2_seq.append(cur.val)
            if cur.right:
                in_order(cur.right, flag)
        if root1:
            in_order(root1, 1)
        if root2:
            in_order(root2, 2)

        while root1_seq or root2_seq:
            if not root1_seq:
                answer += list(root2_seq)
                break
            elif not root2_seq:
                answer += list(root1_seq)
                break

            head1, head2 = root1_seq[0], root2_seq[0]
            if head1 >= head2:
                answer.append(root2_seq.popleft())
            else:
                answer.append(root1_seq.popleft())

        return answer

이진 탐색 트리이기 때문에 중위 순회를 하면 각 트리 별로 오름차순 정렬된 값을 얻을 수 있다.

 

두 트리에서 오름차순 정렬된 값을 담은 list를 0번 인덱스부터 비교하면서 차례대로 넣어주면 된다.

마지막 정렬 과정에서 한쪽 list가 먼저 소진되면 나머지 list를 그냥 붙여주고 반복문을 탈출했다.