[프로그래머스] 등대

문제 요약

n개의 등대와 n - 1개의 뱃길이 있다.

모든 등대는 서로 이동할 수 있도록 연결되어 있으므로 전체 구조는 트리다.

각 뱃길의 양쪽 끝 등대 중 적어도 하나는 켜져 있어야 한다.

모든 뱃길이 이 조건을 만족하도록 켜야 하는 등대 수의 최솟값을 구해야 한다.

핵심 아이디어

각 뱃길을 하나의 간선으로 생각하면 문제의 조건은 다음과 같다.

모든 간선은 선택된 정점을 적어도 하나 포함해야 한다.

그래프의 모든 간선이 선택된 정점과 연결되도록 최소 개수의 정점을 고르는 문제를 최소 정점 커버라고 한다.

이 문제의 그래프는 트리이므로 트리 DP로 최소 정점 커버를 구할 수 있다.

각 등대마다 다음 두 상태를 관리한다.

dp[node][0] = node 등대를 끈 경우,
              node의 서브트리에서 켜야 하는 최소 등대 수

dp[node][1] = node 등대를 켠 경우,
              node의 서브트리에서 켜야 하는 최소 등대 수

부모 등대를 켜는지 끄는지에 따라 자식 등대의 선택 가능 여부가 달라진다.

DP 점화식

현재 등대를 끄는 경우

현재 등대와 자식 등대 사이의 뱃길을 안전하게 만들려면 모든 자식 등대를 반드시 켜야 한다.

현재 등대: OFF
자식 등대: 반드시 ON

따라서 점화식은 다음과 같다.

dp[node][0] += dp[child][1]

모든 자식에 대해 더하면 다음과 같다.

dp[node][0] = sum(dp[child][1])

현재 등대를 켜는 경우

현재 등대가 켜져 있다면 현재 등대와 자식 사이의 뱃길은 이미 조건을 만족한다.

따라서 자식 등대는 켜도 되고 꺼도 된다.

두 상태 중 더 작은 값을 선택한다.

dp[node][1] += min(
    dp[child][0],
    dp[child][1]
)

현재 등대 자신이 켜져 있으므로 초기값은 1이다.

dp[node][1]
= 1 + sum(min(dp[child][0], dp[child][1]))

리프 노드의 초기값

자식이 없는 리프 노드를 살펴보자.

리프 노드를 끄는 경우에는 해당 서브트리에서 켜는 등대가 없다.

dp[leaf][0] = 0

리프 노드를 켜는 경우에는 자기 자신 하나를 켠다.

dp[leaf][1] = 1

따라서 모든 노드의 초기값을 다음과 같이 설정할 수 있다.

dp = [[0, 1] for _ in range(n + 1)]

반복문으로 트리 순회하기

자식의 DP 값을 먼저 계산한 뒤 부모의 값을 계산해야 한다.

즉, 리프 노드에서 루트 방향으로 처리하는 후위 순회가 필요하다.

재귀 DFS를 사용할 수도 있지만 n이 크면 파이썬의 재귀 깊이 제한을 넘을 수 있다.

따라서 스택을 이용해 방문 순서를 만든 뒤 역순으로 처리한다.

부모와 방문 순서 기록

1번 등대를 임의의 루트로 정한다.

parent = [0] * (n + 1)
parent[1] = -1

order = []
stack = [1]

스택으로 트리를 순회하면서 각 노드의 부모와 방문 순서를 저장한다.

while stack:
    node = stack.pop()
    order.append(node)

    for neighbor in graph[node]:
        if neighbor == parent[node]:
            continue

        parent[neighbor] = node
        stack.append(neighbor)

방문 순서를 뒤집어 DP 계산

일반적인 DFS 방문 순서는 부모가 자식보다 먼저 등장한다.

이를 역순으로 확인하면 자식을 부모보다 먼저 처리할 수 있다.

for node in reversed(order):

각 노드에서는 부모 방향을 제외한 이웃만 자식으로 처리한다.

풀이 과정

1. 인접 리스트 생성

각 뱃길은 양방향으로 이동할 수 있으므로 양쪽 등대의 인접 리스트에 모두 추가한다.

graph = [[] for _ in range(n + 1)]

for lighthouse_a, lighthouse_b in lighthouse:
    graph[lighthouse_a].append(lighthouse_b)
    graph[lighthouse_b].append(lighthouse_a)

2. 트리의 부모 관계 생성

1번 등대를 루트로 정하고 반복문 기반 DFS를 수행한다.

parent = [0] * (n + 1)
parent[1] = -1

order = []
stack = [1]

순회하면서 각 등대의 부모와 방문 순서를 기록한다.

3. DP 배열 초기화

dp = [[0, 1] for _ in range(n + 1)]

각 등대를 끄는 경우는 0, 켜는 경우는 자기 자신을 포함해 1로 시작한다.

4. 자식부터 DP 계산

for node in reversed(order):

현재 등대를 끄는 경우에는 자식을 반드시 켠다.

dp[node][0] += dp[child][1]

현재 등대를 켜는 경우에는 자식의 두 상태 중 더 작은 값을 선택한다.

dp[node][1] += min(dp[child][0], dp[child][1])

5. 루트의 두 상태 중 최솟값 반환

루트는 부모가 없으므로 켜도 되고 꺼도 된다.

return min(dp[1][0], dp[1][1])

Python 코드

def solution(n, lighthouse):
    graph = [[] for _ in range(n + 1)]

    for lighthouse_a, lighthouse_b in lighthouse:
        graph[lighthouse_a].append(lighthouse_b)
        graph[lighthouse_b].append(lighthouse_a)

    # 1번 등대를 루트로 삼아 부모와 방문 순서를 기록한다.
    parent = [0] * (n + 1)
    parent[1] = -1

    order = []
    stack = [1]

    while stack:
        node = stack.pop()
        order.append(node)

        for neighbor in graph[node]:
            if neighbor == parent[node]:
                continue

            parent[neighbor] = node
            stack.append(neighbor)

    # dp[node][0]: node를 끈 경우
    # dp[node][1]: node를 켠 경우
    dp = [[0, 1] for _ in range(n + 1)]

    # 자식의 값을 먼저 계산하기 위해 방문 순서를 뒤집는다.
    for node in reversed(order):
        for child in graph[node]:
            if parent[child] != node:
                continue

            # 현재 등대를 끄면 자식 등대는 반드시 켜야 한다.
            dp[node][0] += dp[child][1]

            # 현재 등대를 켜면 자식 등대는 켜거나 끌 수 있다.
            dp[node][1] += min(
                dp[child][0],
                dp[child][1]
            )

    return min(dp[1][0], dp[1][1])

코드 설명

인접 리스트

graph = [[] for _ in range(n + 1)]

등대 번호가 1부터 n까지이므로 크기가 n + 1인 배열을 사용한다.

graph[node]에는 해당 등대와 뱃길로 연결된 모든 등대 번호가 저장된다.

루트 설정

parent[1] = -1

1번 등대를 임의의 루트로 정한다.

트리는 어느 노드를 루트로 선택해도 최종 정답이 달라지지 않는다.

루트의 부모를 -1로 설정하면 다른 노드와 구분할 수 있다.

부모 방향 제외

if neighbor == parent[node]:
    continue

뱃길은 양방향으로 인접 리스트에 저장되어 있다.

부모 등대로 다시 이동하면 같은 노드를 반복해서 방문하게 되므로 부모 방향은 제외한다.

자식 판별

if parent[child] != node:
    continue

DP를 계산할 때 현재 노드와 연결된 이웃 중 현재 노드의 실제 자식만 처리한다.

부모를 자식으로 잘못 처리하면 같은 간선의 값이 중복으로 반영된다.

현재 등대를 끄는 경우

dp[node][0] += dp[child][1]

현재 등대가 꺼져 있다면 현재 등대와 자식 등대를 연결하는 뱃길의 양쪽 중 자식 등대가 반드시 켜져 있어야 한다.

현재 등대를 켜는 경우

dp[node][1] += min(
    dp[child][0],
    dp[child][1]
)

현재 등대가 켜져 있으므로 자식과 연결된 뱃길의 조건은 이미 만족한다.

자식 등대는 켜거나 끌 수 있으며, 해당 서브트리에서 더 적은 등대를 사용하는 상태를 선택한다.

예시

다음과 같은 단순한 트리를 살펴보자.

    1
   / \
  2   3

2번과 3번 등대는 리프 노드다.

dp[2] = [0, 1]
dp[3] = [0, 1]

1번 등대를 끄는 경우에는 2번과 3번을 모두 켜야 한다.

dp[1][0]
= dp[2][1] + dp[3][1]
= 1 + 1
= 2

1번 등대를 켜는 경우에는 2번과 3번을 모두 끌 수 있다.

dp[1][1]
= 1 + min(0, 1) + min(0, 1)
= 1

따라서 최소로 켜야 하는 등대 수는 다음과 같다.

min(2, 1) = 1

1번 등대만 켜면 두 뱃길 모두 적어도 한쪽 끝의 등대가 켜진 상태가 된다.

시간 복잡도

등대의 수를 N이라고 하자.

트리에는 N - 1개의 뱃길이 존재한다.

인접 리스트 생성, 부모 관계 생성, DP 계산 과정에서 각 노드와 간선을 상수 횟수만큼 확인한다.

O(N)

N이 커도 모든 경우를 직접 선택해보지 않기 때문에 효율적으로 처리할 수 있다.

공간 복잡도

인접 리스트, 부모 배열, 방문 순서, DP 배열을 저장한다.

O(N)

정리

이 문제는 트리에서 모든 간선을 덮는 최소 정점 집합을 구하는 최소 정점 커버 문제다.

풀이 흐름은 다음과 같다.

뱃길 정보를 인접 리스트로 구성
1번 등대를 루트로 부모 관계 생성
등대를 켠 상태와 끈 상태로 DP 정의
방문 순서를 역순으로 확인해 자식부터 계산
현재 등대를 끄면 모든 자식을 켬
현재 등대를 켜면 자식의 두 상태 중 최솟값 선택
루트의 두 상태 중 최솟값 반환

현재 등대를 끄는 경우 자식 등대는 반드시 켜야 한다는 상태 전이와, 재귀 깊이 문제를 피하기 위해 반복문으로 후위 순서를 만드는 것이 핵심이다.