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

파이썬으로 트리 색칠하기: 인접 노드가 같은 색을 갖지 않도록 할 수 있는지 확인하는 방법

각 노드의 값이 그 노드의 색을 나타내는 이진 트리가 있다고 가정해 보겠습니다. 트리에는 최대 2가지 색만 존재합니다. 이때 노드들의 색을 원하는 만큼 자유롭게 교환하여 연결된(인접한) 두 노드가 같은 색을 갖지 않도록 만들 수 있는지 확인해야 합니다.

예를 들어 입력 트리가 다음과 같다면,

파이썬으로 트리 색칠하기: 인접 노드가 같은 색을 갖지 않도록 할 수 있는지 확인하는 방법

출력은 True입니다. 아래 그림처럼 색을 재배치하면 인접한 노드가 서로 다른 색을 갖도록 만들 수 있기 때문입니다.

파이썬으로 트리 색칠하기: 인접 노드가 같은 색을 갖지 않도록 할 수 있는지 확인하는 방법

해결 접근 방법

이 문제는 깊이 우선 탐색(DFS)을 한 번만 수행하면 해결할 수 있습니다. 알고리즘의 단계는 다음과 같습니다.

  • colors := 색상별 노드 개수를 저장하는 빈 맵
  • prop := 깊이의 홀짝성(flag)별 노드 개수를 저장하는 빈 맵
  • dfs(node, flag) 함수를 정의합니다.
  • node가 null이면 그대로 반환합니다.
  • colors[노드의 값] := colors[노드의 값] + 1
  • prop[flag] := prop[flag] + 1
  • dfs(노드의 왼쪽 자식, flag 반전)
  • dfs(노드의 오른쪽 자식, flag 반전)
  • 메인 메소드에서는 dfs(root, True)를 호출한 뒤, colors와 prop의 값들이 집합 기준으로 동일하면 true, 그렇지 않으면 false를 반환합니다.

왜 이 방법이 작동할까?

트리는 사이클이 없는 이분 그래프(bipartite graph)입니다. 따라서 노드를 깊이의 홀짝에 따라 두 그룹으로 나누면, 같은 그룹에 속한 노드끼리는 절대 인접하지 않습니다. 즉, 짝수 깊이의 노드는 모두 한 색으로, 홀수 깊이의 노드는 모두 다른 색으로 칠하면 인접한 노드가 항상 다른 색을 갖게 됩니다.

색을 자유롭게 교환할 수 있다는 것은 어떤 노드에 어떤 색 값을 배정할지 임의로 정할 수 있다는 의미입니다. 결국 "색상별 노드 개수의 집합"과 "깊이 홀짝별 노드 개수의 집합"이 일치하기만 하면 유효한 색칠이 항상 존재하며, 이것이 코드의 마지막 비교 조건입니다.

구현 예제

from collections import defaultdict
class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.val = data
        self.left = left
        self.right = right

class Solution:
    def solve(self, root):
        colors = defaultdict(int)
        prop = defaultdict(int)

        def dfs(node, flag=True):
            if not node:
                return
            colors[node.val] += 1
            prop[flag] += 1
            dfs(node.left, not flag)
            dfs(node.right, not flag)

        dfs(root)
        return set(colors.values()) == set(prop.values())

ob = Solution()
root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)
print(ob.solve(root))

입력

root = TreeNode(2)
root.left = TreeNode(2)
root.right = TreeNode(2)
root.right.left = TreeNode(1)
root.right.right = TreeNode(1)
root.right.left.right = TreeNode(1)

출력

True