[코드트리] 코드트리 오마카세 - Python

문제 링크

https://www.codetree.ai/training-field/frequent-problems/problems/codetree-omakase/submissions?page=1&pageSize=20 

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

풀이 

해설에서 아이디어를 참고해서 풀었고, heapq를 사용했다.

각 명령을 실시간으로 처리하는 것이 아니라, 모든 명령들에 대해서 필요한 값들을 미리 저장해두는 작업이 필요하다. 

 

필요한 변수는 아래와 같다. 

l, q = map(int,input().split())
# 전체 쿼리 저장 
query = []
# 사람 별 도착 시간 저장 
entry_time = {}
# 사람 별 위치 저장 
position = {}
# 사람 별 떠나는 시간 저장 
exit_time = defaultdict(int)
# 사람 이름 저장 
names = set()
# 사람 별 쿼리 저장 
p_queries = defaultdict(list)

 

1. 명령어 저장 

heapq를 사용하여, 각 명령을 들어온 시간, 명령 정보 순으로 저장한다. 

들어온 시간이 같다면, 

초밥 만들기 -> 초밥 먹기 -> 새로운 손님 입장하기 -> 손님 떠나기 순으로 처리하고 값을 출력하기 위해서이다. 

따라서 각 명령에 대해 query에 필요한 값을 저장하고, 

초밥 만들기 명령이라면, 손님 별 쿼리에도 추가한다. 

손님 입장하기 명령이라면, 손님 이름을 추가하고 입장 시간과 입장 위치도 저장한다. 

for _ in range(q):
    data = list(map(str,input().split()))
    q_type = int(data[0])
    t, x, n = -1 ,-1 ,-1
    name = ""
    # 스시 들어옴 
    if q_type == 100 :
        # 들어온 시간, 위치, 사람 이름 
        t, x, name = data[1:]
        t = int(t)
        x = int(x)
    if q_type == 200 :
        # 들어온 시간, 위치, 사람 이름, 초밥 개수 
        t, x, name, n = data[1:]
        t = int(t)
        x = int(x)
        n = int(n)
    if q_type == 300 :
        # 사진 촬영 
        t = int(data[1])

    # t와 cmd 순서대로 
    # t가 같다면 초밥 만들고 -> 즉시 초밥 사라지고 -> 새로운 손님 입장하고 -> 손님 사라지고 -> 사진 촬영 순으로
    heapq.heappush(query, (t, q_type, x, name, n))

    if q_type == 100 :
        p_queries[name].append((t, q_type, x, name,n))
    if q_type == 200 :
        names.add(name)
        entry_time[name] = t 
        position[name] = x

 

2. 각 손님 별 떠나는 시간 구하기 

아이디어는 다음과 같다. 

1) 초밥이 손님보다 먼저 만들어 진 경우 

손님이 입장한 시간 때의 초밥 위치를 구해서 손님과 만나기 위해서는 몇 초간 더 이동해야 하는지 구해, 

최종적으로 만나는 시간을 구한다.

손님이 입장한 시간 때의 초밥 위치는 (현재 위치 + (초밥 만들어진 시간 - 손님 입장 시간)) % 의자 개수 이다. 

최종적으로 만나는 시간은 아래 코드와 같다.  

# 초밥과 손님이 만나는 시간 
match_time = 0 

# 초밥이 사람이 오기 전에 주어짐 
if t < entry_time[name]:
    t_diff = entry_time[name] - t 
    # 사람이 들어왔을 때 스시 위치 
    s_pos = (x + t_diff) % l
    match_time = entry_time[name]
    if position[name] > s_pos :
        add_time = position[name] - s_pos
        match_time += add_time
    elif s_pos > position[name]:
        add_time = l - (s_pos - position[name])
        match_time += add_time
    # 몇 초가 더 지나야 만나는지 
    add_time = (position[name] - s_pos + l) % l
    match_time = entry_time[name] + add_time

2) 초밥이 손님이 입장한 이후에 만들어진 경우 

초밥이 만들어진 위치에서 손님까지가기 위해 몇 초간 더 이동해야 하는지 구해,

최종적으로 만나는 시간을 구한다. 

match_time = t
if position[name] > x :
    add_time = position[name] - x
    match_time += add_time
elif x > position[name]:
    add_time = l - (x - position[name])
    match_time += add_time

손님 별 초밥을 먹는 가장 늦은 시간을 갱신하고, 초밥이 사라지는 명령을 추가한다. 

여기서 같은 시간일 경우 초밥이 만들어지고 나서 초밥이 사라지기 때문에 명령어는 100과 200 사이인 111 로 설정 

for name in names :

    for t, cmd, x, p_n, n in p_queries[name]:
        #print(t, cmd, x, p_n, n, entry_time[name], position[name])
        # 초밥과 손님이 만나는 시간 
        match_time = 0 

        # 초밥이 사람이 오기 전에 주어짐 
        if t < entry_time[name]:
            t_diff = entry_time[name] - t 
            # 사람이 들어왔을 때 스시 위치 
            s_pos = (x + t_diff) % l
            match_time = entry_time[name]
            if position[name] > s_pos :
                add_time = position[name] - s_pos
                match_time += add_time
            elif s_pos > position[name]:
                add_time = l - (s_pos - position[name])
                match_time += add_time
            # 몇 초가 더 지나야 만나는지 
            add_time = (position[name] - s_pos + l) % l
            match_time = entry_time[name] + add_time
        else :
            add_time = (position[name] - x + l) % l
            match_time = t + add_time 
        
        # 초밥 먹는 가장 늦은 시간 갱신
        exit_time[name] = max(exit_time[name], match_time)

        # 초밥 사라지는 쿼리 추가 
        heapq.heappush(query, (match_time, 111, -1, name, -1))

 

3. 손님 떠나는 명령 추가 

손님 별 초밥을 먹는 가장 늦은 시간이 손님이 떠나는 시간이므로 2번에서 구한 값을 바탕으로 

시간 별 손님이 떠나는 명령어를 추가한다. 

해당 명령은 200와 300사이의 222로 설정한다.

for name in names:
    heapq.heappush(query, (exit_time[name], 222, -1, name, -1))

 

4. 시점 별 초밥 수, 손님 수 계산

p_num = 0
s_num = 0 

while query:
    data = heapq.heappop(query)
    cmd = data[1]

    if cmd == 100:
        s_num += 1 
    elif cmd == 111 :
        s_num -=1 
    elif cmd == 200:
        p_num += 1 
    elif cmd == 222 :
        p_num -= 1 
    else :
        print(p_num, s_num)

 

전체 코드는 아래와 같다.

#-*- coding: utf-8 -*-
import heapq 
from collections import defaultdict

l, q = map(int,input().split())
# 전체 쿼리 저장 
query = []
# 사람 별 도착 시간 저장 
entry_time = {}
# 사람 별 위치 저장 
position = {}
# 사람 별 떠나는 시간 저장 
exit_time = defaultdict(int)
# 사람 이름 저장 
names = set()
# 사람 별 쿼리 저장 
p_queries = defaultdict(list)

for _ in range(q):
    data = list(map(str,input().split()))
    q_type = int(data[0])
    t, x, n = -1 ,-1 ,-1
    name = ""
    # 스시 들어옴 
    if q_type == 100 :
        # 들어온 시간, 위치, 사람 이름 
        t, x, name = data[1:]
        t = int(t)
        x = int(x)
    if q_type == 200 :
        # 들어온 시간, 위치, 사람 이름, 초밥 개수 
        t, x, name, n = data[1:]
        t = int(t)
        x = int(x)
        n = int(n)
    if q_type == 300 :
        # 사진 촬영 
        t = int(data[1])

    # t와 cmd 순서대로 
    # t가 같다면 초밥 만들고 -> 즉시 초밥 사라지고 -> 새로운 손님 입장하고 -> 손님 사라지고 -> 사진 촬영 순으로
    heapq.heappush(query, (t, q_type, x, name, n))

    if q_type == 100 :
        p_queries[name].append((t, q_type, x, name,n))
    if q_type == 200 :
        names.add(name)
        entry_time[name] = t 
        position[name] = x


for name in names :

    for t, cmd, x, p_n, n in p_queries[name]:
        #print(t, cmd, x, p_n, n, entry_time[name], position[name])
        # 초밥과 손님이 만나는 시간 
        match_time = 0 

        # 초밥이 사람이 오기 전에 주어짐 
        if t < entry_time[name]:
            t_diff = entry_time[name] - t 
            # 사람이 들어왔을 때 스시 위치 
            s_pos = (x + t_diff) % l
            match_time = entry_time[name]
            if position[name] > s_pos :
                add_time = position[name] - s_pos
                match_time += add_time
            elif s_pos > position[name]:
                add_time = l - (s_pos - position[name])
                match_time += add_time
        else :
            match_time = t
            if position[name] > x :
                add_time = position[name] - x
                match_time += add_time
            elif x > position[name]:
                add_time = l - (x - position[name])
                match_time += add_time
        
        # 초밥 먹는 가장 늦은 시간 갱신
        exit_time[name] = max(exit_time[name], match_time)

        # 초밥 사라지는 쿼리 추가 
        heapq.heappush(query, (match_time, 111, -1, name, -1))

# 사람 떠나는 쿼리 추가 
for name in names:
    heapq.heappush(query, (exit_time[name], 222, -1, name, -1))

p_num = 0
s_num = 0 

while query:
    data = heapq.heappop(query)
    cmd = data[1]

    if cmd == 100:
        s_num += 1 
    elif cmd == 111 :
        s_num -=1 
    elif cmd == 200:
        p_num += 1 
    elif cmd == 222 :
        p_num -= 1 
    else :
        print(p_num, s_num)