백준 1620번 나는야 포켓몬 마스터 이다솜 (Python)

문제

포켓몬의 수 N, 내가 맞춰야하는 문제의 수 M.

N개의 줄에 포켓몬의 이름이 입력으로 들어옴. 그 다음 줄부터 M개의 줄에 맞춰야하는 문제가 입력으로 들어옴.

숫자로 들어오면 포켓몬 번호에 해당하는 문자 출력. 알파벳으로 들어오면 포켓몬 번호 출력.

풀이

문제는 쉬운것 같은데 정답 비율이 낮아서 '함정이 있는건가..' 하는 생각이 들었다. 혹시나가 역시나.

처음에는 단순히 python의 내장 라이브러리 함수들을 이용해서 출력하는 방식으로 풀었다. (.index() or .get())

아니나 다를까 계속해서 시간초과 ㅠㅠㅠㅠㅠ

결국에는 이진 탐색을 이용해서 풀고 제출했다. 하지만 내가봐도 너무 지저분하다..

import sys
input = sys.stdin.readline
import operator
 
def findStr(x) :  
    
    s = 0
    l = n-1
 
    while True :             
        ptr = (l + s) // 2        
        if s_pdic[ptr][1] == x :
            break
        elif s_pdic[ptr][1] < x :
            s = ptr + 1
        else :
            l = ptr - 1           
    return ptr
 
n, m = map(int, input().split())
 
pdic = {}
 
for i in range(n) :
    pdic[i] = input().rstrip()
 
s_pdic = sorted(pdic.items(), key=operator.itemgetter(1))
 
for i in range(m) :
 
    tmp = input().rstrip()
    
    if tmp.isdigit() :
        print(pdic.get(int(tmp)-1))
    else :            
        print(s_pdic[int(findStr(tmp))][0]+1)

맞은 사람들의 코드를 보는데 ㅠㅠ 이렇게 쉽게 풀었다니 하는 코드가 몇 있었다.

가장 맘에 들었던 코드를 수정해보았다. (사실 변수 이름만 수정함)

포켓몬 이름을 저장할 list와 포켓몬 이름과 번호를 key-value 형태로 저장하는 dictionary를 따로 두는 것이다.

이런 방식으로 접근하면 몇번째 인덱스인지 탐색할 필요없이 바로 원하는 값을 알아낼 수 있다. 왜 이생각을 못했지??

분발하자..

import sys
input = sys.stdin.readline
 
n, m = map(int, input().split())
 
pkmn = [] # 포켓몬 이름을 저장할 list
pkmn_dic = {} # 포켓몬 이름에 따른 번호를 저장할 dictionary
 
for i in range(n) :    
    pk = input().rstrip()
    pkmn.append(pk)
    pkmn_dic[pk] = i + 1
 
for _ in range(m):
 
    query = input().rstrip()
 
    # query가 숫자이면 list에서 조회 후 출력
    if query.isdigit() :
        print(pkmn[int(query)-1])
    # query가 문자열이면 dictionary에서 조회 후 출력
    else :
        print(pkmn_dic[query])

문제 주소 : https://www.acmicpc.net/problem/1620