[BOJ] 나는야 포켓몬 마스터 이다솜(1620) - PYTHON

포켓몬의 수와 위치를 찾을 개수가 입력으로 들어온다. 이후 포켓몬 이름을 차례대로 입력받은 뒤에, 위치를 찾을 포켓몬 이름 혹은 이름을 찾을 위치(숫자)를 입력받는다.. 포켓몬 이름은 알파벳으로 구성돼있다. 출력할 포켓몬 이름이 들어올 때 숫자가 입력된다면, 해당 숫자 번째로 들어온 포켓몬의 이름을 출력해야 한다. 이름이 입력으로 주어지면 해당 포켓몬의 위치를 출력한다. 

시간 복잡도 : O(1)

1620번: 나는야 포켓몬 마스터 이다솜

---

풀이과정

자료 갯수가 최대 10만 개이므로, 입력이 들어올 때마다 배열을 순회하면 O(N^2)의 시간이 걸리며, 시간 초과가 발생한다. 입력으로 들어오는 포켓몬은 이름이 모두 다르다는 특징이 있으므로 HashMap을 통해 이름을 저장한다. 이때 key는 포켓몬 이름, value는 들어온 순서를 저장한다. 해쉬 맵을 통해 저장하면 key인 포켓몬 이름으로 값을 찾을 때 O(1)의 시간이 걸린다.

한편 숫자가 입력으로 들어오면 해당 순번으로 들어온 포켓몬 이름을 반환해야 하므로 이때는 1차원 배열을 통해 입력받은 포켓몬 이름을 순서대로 저장한 뒤에 인덱싱을 통해 포켓몬 이름을 반환한다. 이때도 역시 O(1)의 시간이 걸린다.

주의할점  

입력을 받을 때 문자열이 최대 10만 개 들어오므로, input함수를 그대로 쓰면 입력이 느리게 받아진다. 따라서 sys모듈을 import 해서 속도를 향상한다. 이때 sys모듈을 사용하면 문자열을 입력받을 때 마지막 개행 문자가 같이 입력받아지므로 rstrip 함수를 통해 공백 문자를 제거한다. 공백을 제거하지 않으면 나중에 찾을 포켓몬 이름 or 찾을 포켓몬 위치 입력 값을 구분할 때 숫자 값을 제대로 읽어오지 못해 오류가 생긴다.

---

코드

import sys
 
input = sys.stdin.readline
N, M = map(int, input().split())
 
pdict = {}
parr = []
for i in range(N):
    s = input().rstrip()
    pdict[s] = i + 1
    parr.append(s)
for j in range(M):
    s = input().rstrip()
    if s.isnumeric():
        print(parr[int(s) - 1])
    else:
        print(pdict[s])

주석 추가

import sys
 
input = sys.stdin.readline  # sys모듈을 통해 입력 속도 향상
N, M = map(int, input().split())  # 포켓몬의 수, 위치 or 이름을 찾을 포켓몬 수
 
pdict = {}  # 포켓몬 사전 (key:포켓몬 이름 value: 입력 순서)
parr = []  # 포켓몬 배열 
for i in range(N):
    s = input().rstrip()  # 입력받고 뒤의 공백 제거
    pdict[s] = i + 1  # 포켓몬이 몇번째로 들어왔는 지 저장
    parr.append(s)  # 해당 위치에 들어온 포켓몬 이름을 저장
for j in range(M):
    s = input().rstrip()  # 찾을 포켓몬 입력
    if s.isnumeric():  # 숫자로 입력이 들어왔으면(포켓몬 이름을 찾아야한다면)
        print(parr[int(s) - 1])  # 배열을 통해 이름 찾기
    else:
        print(pdict[s])  # 사전을 통해 포켓몬 위치 찾기
 

---

채점결과

처음에 배열 인덱싱을 잘못해서 틀렸고, 이후에는 input을 그냥 받아 시간 초과가 발생했다. 그래서 sys모듈로 input을 받았더니 뒤의 공백이 같이 들어가 14번째에서 숫자 입력을 제대로 받지 못해 keyError가 발생했다. 좀 더 꼼꼼히 보고 제출해야겠다.