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

파이썬으로 이진 트리에서 길이가 k인 고유 경로 개수 구하기


문제 설명

서로 다른 고유한 값을 가진 이진 트리와 정수 k가 주어졌다고 가정해 보겠습니다. 이때 트리 안에서 길이가 k인 고유한 경로의 개수를 구해야 합니다. 경로는 부모 노드에서 자식 노드 방향으로 내려갈 수도 있고, 반대로 자식 노드에서 부모 노드 방향으로 거슬러 올라갈 수도 있습니다. 두 경로를 비교했을 때 한쪽에만 포함된 노드가 하나라도 존재한다면, 그 두 경로는 서로 다른 경로로 간주합니다.

예를 들어 다음과 같은 트리가 주어지고,

파이썬으로 이진 트리에서 길이가 k인 고유 경로 개수 구하기

k = 3이라면 출력은 4가 됩니다. 해당하는 경로가 [12,8,3], [12,8,10], [8,12,15], [3,8,10]의 네 가지이기 때문입니다.

해결 전략

이 문제는 깊이 우선 탐색(DFS)을 활용하면 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 노드를 지나는 길이별 경로의 개수를 배열로 관리하고, 왼쪽 자식과 오른쪽 자식의 결과를 조합하여 정답을 누적하는 것입니다.

구체적인 알고리즘은 다음과 같습니다 −

  • dfs() 함수를 정의합니다. 이 함수는 노드를 인자로 받습니다.

    • 노드가 null이라면, 첫 번째 요소가 1이고 나머지 k−1개가 0인 리스트를 반환합니다.

    • left := dfs(노드의 왼쪽 자식)

    • right := dfs(노드의 오른쪽 자식)

    • i를 0부터 K까지 순회하며 다음을 수행합니다.

      • ans := ans + left[i] * right[K − 1 − i]

    • 크기가 K인 0으로 초기화된 리스트 res를 생성합니다.

    • res[0] := 1, res[1] := 1로 설정합니다.

    • i를 1부터 K−1까지 순회하며 다음을 수행합니다.

      • res[i + 1] := res[i + 1] + left[i]

      • res[i + 1] := res[i + 1] + right[i]

    • res를 반환합니다.

  • 메인 메소드에서는 다음을 수행합니다 −

  • ans := 0으로 초기화합니다.

  • dfs(root)를 호출합니다.

  • ans를 반환합니다.


여기서 left[i]와 right[K − 1 − i]를 곱하는 과정은 "왼쪽에서 길이 i로 도달한 경로"와 "오른쪽에서 남은 길이만큼 뻗어 나가는 경로"를 연결해 하나의 k 길이 경로를 만드는 모든 경우의 수를 세는 작업입니다.

예제 코드

다음 구현을 통해 더 잘 이해할 수 있습니다 −

class TreeNode:
   def __init__(self, data, left = None, right = None):
      self.data = data
      self.left = left
      self.right = right
class Solution:
   def solve(self, root, K):
      def dfs(node):
         if not node:
            return [1] + [0] * (K-1)
         left = dfs(node.left)
         right = dfs(node.right)
         for i in range(K):
            self.ans += left[i] * right[K - 1 - i]
         res = [0] * K
         res[0] = res[1] = 1
         for i in range(1, K - 1):
            res[i + 1] += left[i]
            res[i + 1] += right[i]
         return res
      self.ans = 0
      dfs(root)
      return self.ans
ob = Solution()
root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
print(ob.solve(root, 3))

입력

root = TreeNode(12)
root.left = TreeNode(8)
root.right = TreeNode(15)
root.left.left = TreeNode(3)
root.left.right = TreeNode(10)
3

출력

4

복잡도 분석

시간 복잡도는 O(n × K)입니다. n은 트리의 노드 수이며, 각 노드마다 크기 K의 배열 연산을 수행하기 때문입니다. 공간 복잡도는 재귀 호출 스택과 각 노드의 결과 배열을 포함하여 O(n)입니다.