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

백준 1620번 나는야 포켓몬 마스터 이다솜 문제 바로가기

문제 이해

문제가 이상하다.
보고 내 눈이 잘못된줄 아라따..........
문제를 보면 그냥 스토리가 웃기다 ㅎㅎㅋㅎㅋㅎㅋㅎ

그래서 스토리 다 건너띄고 문제부터 봤다. ㅎ
문제를 설명하자면 간단하다.
N M을 입력받고 N개의 포켓몬에는 각각의 번호가 있다. 그 번호를 포켓몬의 번호라고 한다.

M개의 문제가 주어지는데 문제는 모두 숫자 혹은 모두 알파벳으로 이뤄져있다. 따라서 숫자가 주어지면 그 번호에 해당하는 포켓몬의 이름을, 이름이 주어지면 그 포켓몬의 번호를 출력하면 되는 문제다.

문제 접근

처음에는 배열로 접근했다!
그냥 입력받고 문제가 숫자인지 알파벳인지 확인 하고 숫자일 경우에는 list[숫자]로 포켓몬 이름을 출력하고 알파벳일 경우에는 index()를 사용하여 그 index를 구하고자 했다!

하지만 이렇게 작성하고 제출하자 시간초과가 발생했다. 아무래도 내장함수 index() 때문에 발생하는 듯 했다.
아래는 시간초과 코드이다.

import sys
N, M = map(int, sys.stdin.readline().split())
poketmon = []
for i in range(N):
    poketmon.append(str(sys.stdin.readline().strip())) # O(N)
for i in range(M): 
    find = sys.stdin.readline().strip()
    if find[0].isalpha():
        print(poketmon.index(find)+1) # O(N^2)
    else:
        print(poketmon[int(find)-1])

따라서 dictionary 를 사용하여 숫자가 key인 dictionary 와 이름이 key인 dictionary를 각각 만들어줘야겠다고 생각했다. 그렇게 되면 바로 key를 입력하면 되므로 시간초과되지 않을 것이라고 생각했다.

문제 풀이

따라서 아래와 같이 2개의 dict를 생성하고 입력받은 순서와, 알파벳을 고려하여 poketmon dict에는 key에 번호를, value에 포켓몬의 이름을 저장했고, number dict에는 key에 포켓몬의 이름을, value에는 포켓몬의 번호를 저장했다.

그리고 문제는 모두 숫자이거나 모두 알파벳이므로 가장 첫자만 확인하면 숫자인지 알파벳인지 확인할 수 있었다.
M번만큼 for문을 돌려 확인한 값에 따라 dict배열을 조회하면 되는 문제였다.

poketmon = dict()
number = dict()
for i in range(1, N+1):
    po = sys.stdin.readline().strip()
    poketmon[i] = po# O(N)
    number[po] = i

for i in range(M): # O(M)
    find = sys.stdin.readline().strip()
    if find[0].isalpha():
        print(number[find]) 
    else:
        print(poketmon[int(find)])

제출 코드

import sys
N, M = map(int, sys.stdin.readline().split())
poketmon = dict()
number = dict()
for i in range(1, N+1):
    po = sys.stdin.readline().strip()
    poketmon[i] = po# O(N)
    number[po] = i

for i in range(M):
    find = sys.stdin.readline().strip()
    if find[0].isalpha():
        print(number[find]) # O(N^2)
    else:
        print(poketmon[int(find)])