트리 문제를 풀다 보면 막히는 순간이 있습니다.

"이 노드를 루트로 한 부분을 어떻게 판정해야 하지?" 싶은 순간인데, 이 문제는 그 흐름을 연습하기 좋습니다.

이번 글에서는 모든 노드 값이 1인 서브트리의 개수를 구하는 문제를 정리해보겠습니다.

문제

이진 트리의 루트가 주어질 때, 모든 노드 값이 1인 서브트리의 개수를 구하면 됩니다.

여기서 서브트리는 특정 노드를 루트로 하고, 그 아래 자식을 전부 포함한 트리를 뜻합니다.

예를 들어 아래와 같은 트리가 있다고 하겠습니다.

        1
       / \
      1   0
     / \   \
    1   1   0

이 트리에서 값이 전부 1인 서브트리는 총 3개입니다.

조건을 만족하는 서브트리는 다음과 같습니다.

  • 맨 아래 왼쪽 1
  • 맨 아래 가운데 1
  • 그 둘을 자식으로 가지는 왼쪽 서브트리 전체

오른쪽은 0이 포함되어 있으므로 제외됩니다.

처음 떠오르는 생각

처음에는 각 노드를 시작점으로 잡고, 그 아래를 전부 훑으면서 "전부 1인가?"를 검사하면 되겠다고 생각할 수 있습니다.

하지만 그렇게 하면 같은 하위 트리를 여러 번 다시 확인하게 됩니다.

즉, 이 문제의 핵심은 매번 처음부터 다시 세는 것이 아니라 아래에서 먼저 판정한 결과를 부모가 받아 쓰는 구조를 만드는 것입니다.

핵심 아이디어

현재 노드를 루트로 하는 서브트리가 전부 1인지 알려면 결국 세 가지만 확인하면 됩니다.

  • 왼쪽 서브트리가 전부 1인지
  • 오른쪽 서브트리가 전부 1인지
  • 현재 노드 값이 1인지

이 세 가지를 확인하면 현재 노드 기준 판정이 끝납니다.

그래서 재귀 함수 dfs(node)를 아래처럼 정의하면 흐름이 단순해집니다.

dfs(node) = node를 루트로 하는 서브트리가 전부 1이면 True, 아니면 False

이렇게 두면 현재 노드에서는:

  1. 왼쪽 자식 결과를 받는다
  2. 오른쪽 자식 결과를 받는다
  3. 둘 다 괜찮고 현재 값도 1이면 count를 증가시킨다
  4. 현재 결과를 부모에게 반환한다

빈 노드는 조건을 깨뜨리지 않는다고 보고 True를 반환하면 리프 노드도 자연스럽게 처리할 수 있습니다.

코드

class Node:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


def count_all_one_subtrees(root):
    count = 0

    def dfs(node):
        nonlocal count

        if node is None:
            return True

        left_ok = dfs(node.left)
        right_ok = dfs(node.right)

        if not left_ok or not right_ok:
            return False

        if node.val != 1:
            return False

        count += 1
        return True

    dfs(root)
    return count

손으로 한 번 판정해 보기

예시 트리에서 왼쪽 아래 리프 노드 1을 보면, 자식이 없고 자기 값도 1이므로 바로 조건을 만족합니다.

1

이런 리프 노드는 전부 1인 서브트리 1개로 셀 수 있습니다.

그 옆 리프 노드 1도 마찬가지입니다.

이제 그 둘의 부모인 왼쪽 1을 봅니다.

    1
   / \
  1   1

왼쪽 결과도 True, 오른쪽 결과도 True, 현재 값도 1이므로 이 서브트리 전체도 조건을 만족합니다.

반면 오른쪽 0 쪽은 아래에 0이 포함되어 있으므로, 현재 노드까지 올라왔을 때 바로 실패가 전파됩니다.

이 문제는 결국 각 노드가 부모에게 아래 한 문장만 알려주면 됩니다.

"나를 루트로 한 서브트리는 전부 1입니다" 또는 "아닙니다"

이 한 문장만 부모에게 계속 전달하면 전체를 한 번에 처리할 수 있습니다.

코드 해설

빈 노드는 통과로 본다

if node is None:
    return True

빈 노드는 비교를 깨뜨리지 않기 때문에 부모 입장에서는 조건을 통과한 것으로 보는 편이 깔끔합니다.

이렇게 해야 리프 노드도 자연스럽게 처리할 수 있습니다.

왼쪽과 오른쪽을 먼저 본다

left_ok = dfs(node.left)
right_ok = dfs(node.right)

현재 노드를 바로 판단하는 것이 아니라, 먼저 자식 서브트리의 상태를 확인합니다.

자식 중 하나라도 실패하면 현재도 실패한다

if not left_ok or not right_ok:
    return False

왼쪽이나 오른쪽 아래에 1이 아닌 값이 하나라도 있다면 현재 서브트리도 전부 1일 수 없습니다.

현재 값도 1이어야 한다

if node.val != 1:
    return False

자식이 모두 괜찮아도 현재 값이 1이 아니면 실패입니다.

여기까지 왔으면 개수 증가

count += 1
return True

현재 노드를 루트로 하는 서브트리가 조건을 만족한다는 뜻이므로 정답 개수를 하나 늘려줍니다.

왜 위에서 아래로 바로 세면 비효율적일까?

처음에는 "각 노드를 시작점으로 잡고, 아래를 다시 훑어서 전부 1인지 확인하면 되겠지"라고 생각할 수 있습니다.

예를 들어 루트에서 한 번 전체를 보고, 그다음 왼쪽 자식을 루트로 다시 보고, 또 그다음 자식에서 다시 보는 식입니다.

이 방식은 틀리지는 않지만, 같은 하위 트리를 여러 번 반복해서 확인하게 됩니다.

그래서 트리 크기가 커질수록 비효율적입니다.

반면 지금 풀이에서는 각 노드가 딱 한 번만 계산됩니다.

  • 자식이 먼저 자기 결과를 만든다
  • 부모는 그 결과를 받아서 자기 판정을 한다
  • 다시 같은 서브트리를 처음부터 보지 않는다

즉 이 문제는 "개수를 세는 문제"이기도 하지만, 동시에 판정 결과를 위로 올려 보내는 DFS 문제이기도 합니다.

실행 예시

root = Node(
    1,
    Node(1, Node(1), Node(1)),
    Node(0, None, Node(0))
)

print(count_all_one_subtrees(root))  # 3

실수하기 쉬운 부분

None을 False로 두면 리프 노드가 깨집니다

if node is None:
    return True

이 부분이 False가 되면, 리프 노드는 왼쪽과 오른쪽 자식이 모두 실패한 것처럼 보이게 됩니다. 그래서 조건을 만족하는 가장 단순한 경우부터 틀어집니다.

현재 값이 1인지 확인하기 전에 자식 결과를 먼저 받아야 합니다

이 문제는 현재 노드만 보는 문제가 아닙니다. 아래 서브트리가 이미 실패했는지를 먼저 알아야 하므로, DFS 순서가 중요합니다.

개수는 "현재 서브트리가 전부 1일 때만" 증가해야 합니다

리프 노드에서 1이라고 해서 무조건 올리면 안 되고, 왼쪽과 오른쪽 결과까지 모두 통과했을 때만 count를 증가시켜야 합니다.

복잡도

  • 시간 복잡도 — O(n)
  • 공간 복잡도 — O(h)

모든 노드를 한 번씩만 방문하므로 시간 복잡도는 O(n)입니다.

추가 공간은 재귀 호출 스택 때문에 트리 높이 h만큼 필요합니다.

정리

  • 이 문제는 현재 서브트리가 전부 1인지 판정하는 문제입니다
  • dfs(node)가 불리언을 반환하게 두면 구조가 단순해집니다
  • 자식 결과를 먼저 본 뒤 현재 값을 확인하면 됩니다
  • 조건을 만족하는 순간 개수를 증가시키면 됩니다

트리 문제를 풀 때는 "현재 노드가 맞는가?"보다 "그 판단을 위해 자식에게 먼저 무엇을 물어봐야 하는가?" 를 먼저 떠올리는 편이 더 잘 풀립니다.

비슷한 트리 DFS 유형을 공부하면서, 조건을 바꿔 다시 구성해본 예제입니다.