[Python 3] BOJ - 3430 "용이 산다"

 # 문제 링크

https://www.acmicpc.net/problem/3430

3430번: 용이 산다

 

 # 풀이

 개인적으로 이 문제의 핵심은 그리디 알고리즘과 우선순위 큐라고 생각한다. 현재 호수를 마실 수 있는 차례라면, 1차적으로는 비지 않은 호수 중 하나를 마실 수 있다. 그런데, 그런 호수들 중에서 앞으로 비가 가장 먼저 올 호수를 비우는 것이 이득임은 자명하다. 따라서, 현재 호수를 마실 수 있다면, 비어 있는 호수들을 확인해야 한다

 그런데 문제에서 호수의 갯수가 백만개, 날씨의 정보 역시 길이가 백만이므로 모든 호수를 보는 $O(N^2)$의 풀이로는 절대로 시간안에 정답을 도출할 수 없다. 그러나 윗 단락에서 언급하였듯, 용이 호수를 마시는 순간이 온다면 필요한 정보는 마실 수 있는, 즉 비어있지 않는 호수 중 비가 오는 시간이 가장 빠른 호수이다. 따라서, 우선순위 큐가 필요하다.

<그림 1. 우선순위 큐를 채우는 방법>

 위 그림 1은 임의의 상황을 나타낸 그림이다. 만일, 현재 i번째 인덱스까지 재앙이 일어나지 않는다고 가정하자. 따라서, 현재 i번째 인덱스에 따라 호수 3이 현재 비어있다는 뜻이다. 그런데, 현재 i번째 인덱스에 의해 호수 3은 채워진다. 이 때 i+1번째 인덱스는 다시 용이 호수 중 하나를 마실 수 있는 순간이다. 그런데 만일 i번째에서 우선순위 큐에 대한 처리가 이루어지지 않는다면, i+1번째 인덱스의 입장에서는 마실 수 있는 후보군을 하나 놓치는 것과 마찬가지이다. 따라서, 현재 상태가 비가 오는 상태이면서 비가 오는 호수가 비어져 있다면, 즉 현재 호수에 비가 와도 재앙이 발생하지 않는 경우라면 현재 호수가 다음에 비가 올 시간을 우선순위 큐에 저장해야 한다. 

 위와 같이 우선순위 큐를 채우게 된다면, 자연스럽게 용이 호수를 마시는 순간에 우선순위 큐에 들어가 있는 원소는 무조건 용이 호수를 마시는 시간보다 늦기 때문에 마실 수 있는 호수가 된다. 또한 비가 와서 현재 호수를 채우기 때문에 비어있는 호수를 용이 마시는 경우도 없다. 이 때, 각 호수를 기준으로 각자의 비오는 시간을 전처리해준다면 $O(NlogN)$만에 정답을 구할 수 있다

 

 # 코드

import sys, heapq

# 입력부
tc = int(sys.stdin.readline())
for _ in range(tc):
    n, m = map(int, sys.stdin.readline().split())
    arr = list(map(int, sys.stdin.readline().split()))
    
    # empty : 현재 호수가 비었는지 아닌지를 저장하는 리스트
    empty = [False] * (n + 1)
    # can : 재앙을 피할 수 있는지 없는지 나타내는 flag 변수
    can = True
    # cache : 현재 호수 번호가 i일 때, 비가 오는 시간을 역순으로 저장하는 2차원 배열
    cache = [[] for _ in range(n + 1)]
    # time : 현재 시간 i에서 비가 올 때 해당하는 호수의 번호
    time = [-1] * m
    for i in range(m - 1, -1, -1):
        if arr[i]:
            time[i] = arr[i]
            cache[arr[i]].append(i)
            
    # 가장 첫 시간을 우선순위 큐에 삽입
    q = []
    for i in range(1, n + 1):
        if cache[i]:
            heapq.heappush(q, cache[i].pop())
            
    # ans : 용이 마시는 호수 번호를 저장하는 리스트
    ans = []
    for i in range(m):
        # 현재 비가 온다면
        if arr[i]:
            # 비가 오는데 호수가 비지 않으면 재앙이 일어남
            if not empty[arr[i]]:
                can = False
                break
            # 비가 오는데 호수가 비었다면 그 다음 비오는 시간을 우선순위 큐에 저장
            empty[arr[i]] = False
            if cache[arr[i]]:
                heapq.heappush(q, cache[arr[i]].pop())
                
        # 비가 오지 않는다면
        else:
            # 마실 수 있는 호수가 있다면 그리디하게 마신다
            if q:
                now = q[0]
                heapq.heappop(q)
                ans.append(time[now])
                empty[time[now]] = True
            # 마실 수 있는 호수가 없으니 0을 저장한다
            else:
                ans.append(0)
                
    # 정답 출력
    if not can:
        print('NO')
    else:
        print('YES')
        print(' '.join(map(str, ans)))