[백준 28219] 2023 정올 1차 주유소

문제 출처 : https://www.acmicpc.net/problem/28219

28219번: 주유소

이 문제는 2023년도 정보 올림피아드 1차 중등부 3번, 고등부 2번 문제 입니다.

문제 이해하기

N개의 마을이 있는데 도로가 N - 1개가 있습니다. 정점의 갯수보다 간선의 갯수가 1개 적다면 트리 모형이 아닌가 고민해봐야 합니다. 여기에 임의의 두 마을에 대해 두 마을을 잇는 경로가 정확히 하나 존재한다고 합니다. 이 말은 트리 형태의 마을이라는 뜻입니다. 임의의 마을로 가는데 경로가 두 개 이상이라면 사이클이 있다는 뜻이고, 그것은 트리가 아닙니다.

트리 형태의 마을에서 주유소를 설치하려면 가장 말단에 위치한 마을부터 거리를 확인하면서 최대한 멀리 지을 수 있을 때 주유소의 최솟값을 구할 수 있습니다. 예를 확인해 보면서 어떻게 주유소를 지어야 하는지 확인해 보겠습니다.

이런 형태의 마을이 있습니다. 거리 4안에 주유소를 한개이상 설치해야 한다고 하겠습니다. 말단 정점인 11, 9, 6, 10들로 부터 부모 정점으로 한칸씩 올라가면서 거리를 확인합니다.

만나는 점이 생기는 2까지 올라오면서 거리를 확인해 보면 11에서 올라온 거리가 3, 9에서 올라온 거리가 2, 6에서 올라온 거리가 1입니다. 가장 긴 두 거리인 3과 2의 합이 5가 되어 4보다 크게 됩니다. 그럼 2에다가 주유소를 지으면 거리 4안에 주유소가 생기게 됩니다. 즉 우리는 여러개의 노드를 타고 올라오는 거리 중에서 가장 긴 두 개의 거리를 기억해야 합니다.

 

2번에 주유소를 설치하고 10번 정점부터 위로 올라오도록 해보겠습니다.

10에서 올라온 거리가 3이고 2에서 올라온 거리가 0이 되어야 1번에 주유소를 설치 안하게 됩니다. 2번에서 올라온 거리가 1이라면 1번도 거리가 4가 되어 주유소를 설치해야 되기 때문입니다. 따라서 우리는 주유소가 설치된 마을의 값을 -1로 하여 거리 계산에 추가되지 않도록 해야 합니다. dfs를 사용하여 그리디로 주유소를 설치하면 되는 문제입니다.

코드 작성하기

그럼 지금까지 알아본 내용을 바탕으로 코드를 작성해 보겠습니다.

입력 받기

import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

mii = lambda : map(int, input().split())

N, K = mii()

arr = [[] for _ in range(N + 1)]
for _ in range(N - 1):
    u, v = mii()
    arr[u].append(v)
    arr[v].append(u)

먼저 말단 노드부터 재귀를 사용할 것이기 때문에 setrecursionlimit() 함수를 사용하여 재귀의 한계치를 올려줍니다. 그리고 여러 입력이 들어오기 때문에 input을 readline으로 바꿔 속도를 빠르게 해줍니다.

다음으로 마을의 수 N과 주유소의 최소 거리 K를 입력 받습니다. 그리고 u, v 값으로 트리를 구성해 줍니다.

초기화 하기

visited = [False] * (N + 1)
dp = [0] * (N + 1)

ans = 0 

dfs를 사용할 예정이기 때문에 visited 리스트를 만들어 방문 여부를 체크합니다. dp는 해당 마을에서 최대 거리를 저장합니다. 그리고 ans에는 주유소의 갯수를 저장할 것입니다.

dfs 함수 만들기

def dfs(node):
    global ans 
    visited[node] = True

    dist1, dist2 = 0, 0
    for next_node in arr[node]:
        if not visited[next_node]:
            dfs(next_node)

            next_dist = dp[next_node] + 1

            if dist1 <= next_dist:
                dist2 = dist1
                dist1 = next_dist
            elif dist2 <= next_dist:
                dist2 = next_dist
    
    if K <= dist1 + dist2:
        ans += 1
        dp[node] = -1
    else:
        dp[node] = dist1

ans값을 global로 선언하여 함수 내에서도 사용 가능하도록 합니다. 마을에 도착하면 먼저 방문 여부를 체크해 줍니다. 노드를 돌아다니면서 가장 긴 두 거리를 찾습니다. next_dist는 이전 마을에서 1 늘린 거리로 이전 마을에서 가장 긴 거리를 찾는 것입니다.

찾은 거리 next_dist와 가장 긴 거리 dist1을 비교합니다. 만약 next_dist가 더 길다면 기존의 가장 긴 거리인 dist1을 두 번째 긴 거리인 dist2로 바꿔주고 next_dist의 거리를 dist1으로 바꿔 줍니다. 만약 가장 긴 거리는 아니지만 두 번째 거리 dist2 보다 긴 거리가 발견된다면 dist2를 두 번째로 긴 거리로 업데이트 합니다.

모든 거리를 계산한 뒤 dist1과 dist2의 합이 K보다 길다면 주유소를 설치합니다. 주유소 설치가 되었기 때문에 ans값을 1 늘려줍니다. 그리고 해당 마을의 거리를 -1로 초기화 합니다. 0으로 하지 않은 이유는 앞에서 설명했듯이 주유소는 거리 계산에 추가되지 않게 하기 위해서 입니다.

만약 계산한 거리가 K보다 작다면 가장 긴 거리 dist1을 마을의 거리로 업데이트 해줍니다.

출력하기

dfs(1)
print(ans)

dfs를 돌려 나온 주유소의 갯수 ans를 출력해 줍니다. 이렇게 작성해서 제출하면 python3에서는 100점이 나오지만 pypy3에서는 마지막에 시간초과가 발생합니다. 따라서 pypy3에서는 재귀를 사용하지 않고 반복문으로 작성해야 합니다. bfs로 작성하면 좋겠지만 밑에서부터 올라오면서 거리를 계산해야 하기 때문에 bfs와 다릅니다.

반복문으로 dfs 구현하기

반복문으로 dfs를 구현해 보겠습니다. 단순 dfs를 구현하는 것은 어렵지 않지만 우리는 올라오면서 계산할 것이 많기 때문에 조금 복잡합니다.

말단 정점까지 진행하기

visited = [False] * (N + 1)
dp = [-1] * (N + 1)
stack = [1]

while que:
    node = stack[-1]

    if not visited[node]:
        visited[node] = True
        for next_node in arr[node]:
            if not visited[next_node]:
                stack.append(next_node)
        continue    

마을의 거리인 dp값을 -1로 초기화 하였습니다. 마지막 정점 계산시 0이 되게 하기위해서 입니다. 마지막 마을의 경우 연결된 마을이 부모 마을 하나가 있습니다. 부모 마을로 인해 거리 계산시 1이 되기 때문에 초깃값을 -1로 하여 마지막 마을의 거리를 0으로 만들어 주기 위해서 입니다.

반복문을 실행하면 모든 노드를 방문할 때까지 stack에 정점이 쌓이게 됩니다. 마지막 정점까지 모두 쌓이도록 만들어 주기 위해서 큐가 아닌 스택을 사용하였습니다.

거리 계산하기

    dist1, dist2 = 0, 0
    for next_node in arr[node]:
        next_dist = dp[next_node] + 1
        if dist1 <= next_dist:
            dist2 = dist1
            dist1 = next_dist
        elif dist2 <= next_dist:
            dist2 = next_dist
    
    if K <= dist1 + dist2:
        ans += 1
        dp[node] = -1
    else:
        dp[node] = dist1

    stack.pop()

print(ans)

거리 계산은 이전과 같습니다. 다른점은 위에 설명하였듯이 마지막 마을의 거리가 0이 되도록 하기 위해서 초깃값을 -1로 했다는 점입니다.

전체 코드

전체 코드를 확인해 보겠습니다.

import sys
input = sys.stdin.readline

mii = lambda : map(int, input().split())

N, K = mii()

arr = [[] for _ in range(N + 1)]
for _ in range(N - 1):
    u, v = mii()
    arr[u].append(v)
    arr[v].append(u)

visited = [False] * (N + 1)
dp = [-1] * (N + 1)

ans = 0
stack = [1]

while stack:
    node = stack[-1]

    if not visited[node]:
        visited[node] = True
        for next_node in arr[node]:
            if not visited[next_node]:
                stack.append(next_node)
        continue

    dist1, dist2 = 0, 0
    for next_node in arr[node]:
        next_dist = dp[next_node] + 1
        if dist1 <= next_dist:
            dist2 = dist1
            dist1 = next_dist
        elif dist2 <= next_dist:
            dist2 = next_dist
    
    if K <= dist1 + dist2:
        ans += 1
        dp[node] = -1
    else:
        dp[node] = dist1

    stack.pop()

print(ans)