Python 알고리즘 인터뷰 - 문제 46. 두 이진 트리 병합

leetcode 617. https://leetcode.com/problems/merge-two-binary-trees/description/

 

문제

두 이진 트리를 병합하라. 중복되는 노드는 값을 합산한다.

예제1

입력: root1 = [1,3,2,5], root2 = [2,1,3,None,4,None,7]

출력: [3,4,5,5,4,None,7]

 

예제2

입력: root1 = [1], root2 = [1,2]

출력: [2, 2]

 

풀이

mergeTree함수를 다음과 같이 정의 한다. 현재 t1, t2를 기준으로 두 노드가 동시에 존재하면 t1.val + t2.val 값으로 노드를 생성하고, 좌측 노드에 mergeTree(t1.left, t2.left)의 반환값을, 우측 노드에 mergeTree(t1.right, t2.right)의 반환값을 재귀적으로 호출하여 연결한다.

 

코드는 다음과 같다.

def mergeTrees(self, t1, t2):
    if t1 and t2:
        node = TreeNode(t1.val + t2.val)
        node.left = self.mergeTrees(t1.left, t2.left)
        node.right = self.mergeTrees(t1.right, t2.right)
        return node

 

해당 코드는 좌측 노드부터 순회를 진행한다. 두 노드가 존재하면 재귀적으로 노드 생성 및 연결을 진행한다.

순회를 진행하는 중 t1, t2중 노드가 존재하지 않으면 둘 중 존재하는 노드를 반환한다. 

반환 코드는 다음과 같다.

def mergeTrees(self, t1, t2):
    if t1 and t2:
        ...
    else:
    	return t1 or t2

좌측 순회가 끝나면 백트래킹이되며 우측으로 순회하면서 merge가 진행된다.

전체 코드

class Solution(object):
    def mergeTrees(self, t1, t2):
        if t1 and t2:
            node = TreeNode(t1.val + t2.val)
            node.left = self.mergeTrees(t1.left, t2.left)
            node.right = self.mergeTrees(t1.right, t2.right)
            return node
        else:
            return t1 or t2