2019 KAKAO 겨울 인턴쉽 호텔 방 배정

https://programmers.co.kr/learn/courses/30/lessons/64063#

코딩테스트 연습 - 호텔 방 배정

 

 

Naive하게 생각하면, 요청이 올때마다 해당 방이 빈방인지 아닌지 check하고, 빈방이 아니라면 방번호를 1씩 늘려가며 빈 방인지 아닌지 확인하는 것이 먼저 떠오를 것이다.

def solution(k, room_number):
    answer = []
    now = dict()
    
    for num in room_number:
        if num not in now:
            now[num] = num + 1
        else:
            while num in now:
                num = now[num]
            now[num] = num+1
        answer.append(num)
        
    return answer

 

하지만, 요청이 총 200,000개 여서 O(n^2) 이 걸리게 되어 TLE에 걸리게 된다.

방번호를 늘려가면서 빈 방을 찾을 때, 찾은 빈 방을 지금까지 확인했던 방과 연결 시켜, 앞으로 해당 방들에 대한 요청이 들어오면 현재 찾은 빈칸을 바로 반환하게 하면 어떨까? 라는 생각이 들었다.

그래서 현재 요청을 처리하기 위해 확인 했던 방들을 저장하고, 빈 방을 찾은 뒤, 저장된 방들의 다음방을 빈 방으로 연결시켜주는 작업을 했다.

다음은 구현한 코드이다.

import sys
sys.setrecursionlimit(500000)

def solution(k, room_number):
    answer = []
    now = dict()
    
    for num in room_number:
        if num not in now:
            now[num] = num + 1
        else:
            stack = []
            while num in now:
                num = now[num]
                stack.append(num)
            now[num] = num+1
            for n in stack:
                now[n] = num+1
        answer.append(num)
        
    return answer

위와 같은 최적화를 하니 통과할 수 있었다.

이런 경로압축은 재귀함수를 통해 구현할 수 있다. Union-Find를 알고 있다면 쉽게 구현할 수 있을 것이다.

다음은 Union-Find로 구현한 코드이다.

import sys
sys.setrecursionlimit(500000)


def solution(k, room_number):
    answer = []
    now = dict()
    
    def find(a):
        if a not in now : return a
        else:
            now[a] = find(now[a])
            return now[a]
    
    def union(a, b):
        a = find(a)
        b = find(b)
        now[a] = b


    for number in room_number:
        n = find(number)
        answer.append(n)
        union(number, n+1)
        
    
    return answer