Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬(Python)으로 이진 트리에서 좋은 리프 노드 쌍의 개수 찾기

문제 설명

이진 트리와 거리 값 d가 하나 주어집니다. 서로 다른 두 리프(잎) 노드로 이루어진 쌍은, 두 노드 사이의 최단 경로 길이가 d보다 작거나 같을 때 '좋은(good) 쌍'으로 정의됩니다.

예를 들어 아래와 같은 이진 트리가 있다고 가정해 보겠습니다.

파이썬(Python)으로 이진 트리에서 좋은 리프 노드 쌍의 개수 찾기

이때 거리 d = 4라면 정답은 2입니다. 그 이유는 (8, 7)과 (5, 6) 두 쌍의 경로 길이가 각각 2로 조건을 만족하지만, (7, 5)나 (8, 6) 같은 나머지 쌍들은 경로 길이가 5로 d = 4보다 크기 때문에 좋은 쌍이 될 수 없습니다.

해결 접근 방법

이 문제는 후위 순회(post-order traversal) 기반의 재귀 함수를 활용해 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 노드에서 자손 리프 노드까지의 거리 정보를 리스트로 관리하고, 왼쪽 서브트리와 오른쪽 서브트리의 리프 노드들을 서로 조합하며 조건을 만족하는 쌍을 세는 것입니다.

구체적인 단계는 다음과 같습니다.

  • 정답 카운터 sol을 0으로 초기화합니다.

  • 루트 노드를 인자로 받는 util() 함수를 정의합니다.

  • 루트가 null이면 빈 리스트를 반환합니다.

  • 루트가 리프 노드라면 [0, 0] 형태의 항목 하나를 담은 배열을 반환합니다. 여기서 두 번째 요소는 현재 노드부터 해당 리프까지의 거리를 의미합니다.

  • 그 외의 경우에는 다음을 수행합니다.

    • 왼쪽 자식에 대한 util() 호출 결과를 l에 저장합니다.

    • 오른쪽 자식에 대한 util() 호출 결과를 r에 저장합니다.

    • l과 r의 모든 항목 n에 대해 n[1] 값을 1씩 증가시킵니다. 현재 노드를 한 단계 거치므로 거리가 1 늘어나기 때문입니다.

    • r의 각 항목 n과 l의 각 항목 n1에 대해 n[1] + n1[1] <= d를 만족하면 sol을 1 증가시킵니다. 즉, 왼쪽과 오른쪽 서브트리에 속한 리프 노드 쌍의 경로 길이가 d 이하인 경우를 셉니다.

    • 병합된 리스트 l + r을 상위 호출로 반환합니다.

  • 메인 메서드에서 util(root)를 호출한 뒤 sol을 반환합니다.

이 알고리즘의 시간 복잡도는 최악의 경우 O(n²)이며, 공간 복잡도는 재귀 스택과 거리 리스트 저장을 위해 O(n)입니다.

아래 예제 구현을 통해 더 잘 이해해 보겠습니다.

예제

class TreeNode:
def __init__(self, val=0, left=None, right=None):
   self.val = val
   self.left = left
   self.right = right
class Solution:
   def __init__(self):
      self.sol = 0
   def solve(self, root, d):
      def util(root):
         if not root:
            return []
         if not root.left and not root.right:
            return [[0, 0]]
         else:
            cur = []
            l = util(root.left)
            r = util(root.right)
            for n in l:
               n[1] += 1
            for n in r:
               n[1] += 1
            for n in r:
               for n1 in l:
                  if n[1] + n1[1] <= d:
                     self.sol += 1
            return l+r
      util(root)
      return self.sol
root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.right = TreeNode(4)
root.left.right.left = TreeNode(8)
root.left.right.right = TreeNode(7)
root.right.left = TreeNode(5)
root.right.right = TreeNode(6)
d = 4
ob = Solution()
print(ob.solve(root, d))

입력

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.right = TreeNode(4)
root.left.right.left = TreeNode(8)
root.left.right.right = TreeNode(7)
root.right.left = TreeNode(5)
root.right.right = TreeNode(6)
d = 4

출력

2

실행 결과로 2가 출력되는데, 이는 경로 길이가 d = 4 이하인 좋은 리프 노드 쌍이 (8, 7)과 (5, 6)으로 총 두 개 존재한다는 의미입니다.