[백준] 19637 IF문 좀 대신 써줘 - python

백준 19637 IF문 좀 대신 써줘 [링크]

문제

게임 개발자인 밀리는 전투력 시스템을 만들어, 캐릭터가 가진 전투력을 기준으로 칭호를 붙여주려고 한다.

예를 들어, 전투력 10,000 이하의 캐릭터는 WEAK, 10,000 초과 그리고 100,000 이하의 캐릭터는 NORMAL, 100,000 초과 그리고 1,000,000 이하의 캐릭터는 STRONG 칭호를 붙여준다고 하자. 이를 IF문으로 작성한다면 아래와 같이 구현할 수 있다.

if power <= 10000
 print WEAK
else if power <= 100000
 print NORMAL
else if power <= 1000000
 print STRONG

혼자서 게임을 개발하느라 매우 바쁜 밀리를 대신해서, 캐릭터의 전투력에 맞는 칭호를 출력하는 프로그램을 작성하자.

입력

첫 번째 줄에는 칭호의 개수 N (1 ≤ N ≤ 105)과 칭호를 출력해야 하는 캐릭터들의 개수 M (1 ≤ M ≤ 105)이 빈칸을 사이에 두고 주어진다. (1 ≤ N, M ≤ 105)

두 번째 줄부터 N개의 줄에 각 칭호의 이름을 나타내는 길이 1 이상, 11 이하의 영어 대문자로만 구성된 문자열과 해당 칭호의 전투력 상한값을 나타내는 109 이하의 음이 아닌 정수가 주어진다. 칭호는 전투력 상한값의 비내림차순으로 주어진다.

N + 2번째 줄부터 M개의 각 줄에는 캐릭터의 전투력을 나타내는 음이 아닌 정수가 주어진다. 해당하는 칭호가 없는 전투력은 입력으로 주어지지 않는다.

출력

M개의 줄에 걸쳐 캐릭터의 전투력에 맞는 칭호를 입력 순서대로 출력한다. 어떤 캐릭터의 전투력으로 출력할 수 있는 칭호가 여러 개인 경우 가장 먼저 입력된 칭호 하나만 출력한다.

---

시간 초과난 경우

시간 대비 입력이 크다는 것을 생각하지 못하고, O(n^2)으로 구현하였다가 시간초과로 실패하였다.

# IF문 좀 대신 써줘
import sys

n, m = map(int, sys.stdin.readline().split(' '))
strong = []
for i in range(n):
    a, b = sys.stdin.readline().split(' ')
    strong.append([a, int(b), i])
    
answer = []

strong.sort(key= lambda x : (x[1], x[2]))

for j in range(m):
    power = int(sys.stdin.readline().rstrip())
    for k in range(n):
        if k == 0:
            if power <= strong[k][1]:
                answer.append(strong[k][0])
        if strong[k-1][1] < power <= strong[k][1]:
            answer.append(strong[k][0])

for h in answer:
    print(h)

시간을 줄이기 위한 아이디어 - 이진탐색

시간을 줄이기 위해 이진탐색을 사용하였다. 전투력를 기준으로 오름차순 배열로 만들 수 있기 때문에 이진탐색을 사용할 수 있다.
그러나 이진탐색을 사용하고도 시간초과가 났었다.

이진탐색 - 시간초과 버전

# IF문 좀 대신 써줘
# 시간초과 나서 이진탐색으로 해야 한다...
import sys

n, m = map(int, sys.stdin.readline().split(' '))

menu = []
nums = []
answer = []

for i in range(n):
   a, b = sys.stdin.readline().split(' ')
   if int(b) not in nums:
       menu.append([a, int(b)])
       nums.append(int(b))

menu.sort(key=lambda x : x[1])

for j in range(m):
   temp = int(sys.stdin.readline().rstrip())
   start = 0
   end = len(menu) - 1
   while start <=  end:
       mid = (start+end) // 2
       if temp > menu[mid][1]:
           start = mid + 1
       else:
           end = mid - 1
   print(menu[start][0])

그 이유는 전투력이 같은 칭호는 걸러주려고 if문을 한 번 더 썼기 때문이다.
(출력할 수 있는 칭호가 여러 개인 경우 가장 먼저 입력된 칭호 하나만 출력한다.) -> 이 조건 때문에!
그러나 이 경우엔 temp > menu[mid][1] 일 경우 end를 재배정하기 때문에 같은 전투력일 때도 가장 앞의 칭호를 가지게 된다.

이진탐색 - 시간 초과 X

# IF문 좀 대신 써줘
# 시간초과 나서 이진탐색으로 해야 한다...
# 이진탐색을 해도 시간초과가 났음. 
# menu를 만들 때 같은 능력치 값이 있으면 그것을 거르는 과정을 거쳤는데
# 이를 거르니까 시간초과가 나지 않음.
import sys

n, m = map(int, sys.stdin.readline().split(' '))

menu = []

for i in range(n):
    a, b = sys.stdin.readline().split(' ')
    menu.append([a, int(b)])

# 이진 탐색을 위해 오름차순으로 정렬
menu.sort(key=lambda x : x[1])

for j in range(m):
    temp = int(sys.stdin.readline().rstrip())
    start = 0
    end = len(menu) - 1
    while start <=  end:
        mid = (start+end) // 2
        if temp > menu[mid][1]:
            start = mid + 1
        else:
            end = mid - 1
    print(menu[start][0])

따라서 이진탐색으로 최종적으로 해결하였다.