[Python/파이썬] 백준 12757번 전설의 JBNU

12757번: 전설의 JBNU

 

문제

전설의 프로그래머 윤준하는 독자적인 데이터베이스 시스템 JBNU(Jeong Bo Neoh Um)를 만들었다.

준하가 생각한 데이터베이스의 기본 골자는 데이터에 접근하기 위한 Key와 그 데이터를 나타내는 Value로 구성되어 있다. 사용자는 Key를 알고 있어야만 원하는 데이터에 접근할 수 있다.

하지만 준하는 건망증이 심해 Key를 매번 잊어버리기 일쑤였다. 따라서 준하는 JBNU를 개조하여 잘못된 Key를 입력하더라도 그 잘못된 Key와 제일 근접한 Key를 찾아주는 메커니즘을 도입하였다.

Key와 Value는 항상 정수로 되어있다. 가장 근접한 Key란 두 수의 차이가 가장 작은 Key를 의미한다. 또한, 정보의 정확성을 위해 두 수의 차이가 K보다 큰 경우는 Key로 인정하지 않기로 하였다.

프로젝트 베끼기의 달인 승균이는 데이터베이스 시간에 JBNU를 모방하기로 했다. 그러나 준하는 전설이기 때문에 그가 만든 프로그램은 찾을 방법이 없었고, 하는 수 없이 같은 조원인 당신에게 맡기려고 한다.

JBNU의 초기 데이터 상태가 주어질 때, 데이터 추가, 수정 및 검색을 지원하는 프로그램을 작성해보자.

입력

첫 줄에는 초기 데이터의 개수인 N(1 ≤ N ≤ 100,000) 과 명령 횟수인 M(1≤M≤100,000), 가장 근접한 Key까지의 거리의 제한인 K (1≤ K ≤10,000)가 주어진다. 

입력의 둘째 줄부터 N개의 줄에는 초기 데이터인 Key와 Value 값이 주어진다. 모든 Key와 Value는 1,000,000,000 이하의 음이 아닌 정수이다. 같은 Key를 갖는 데이터는 없다.

다음 M개의 줄에는 아래와 같은 명령이 주어진다.

• 1 Key Value : 해당 Key와 Value를 가진 데이터를 추가한다. Key가 이미 존재하는 입력은 주어지지 않는다.
• 2 Key Value : 해당 Key로 검색된 데이터를 Value로 변경한다. 조건을 만족하는 유일한 Key가 없는 경우 무시한다.
• 3 Key : 해당 Key로 검색된 데이터를 출력한다. 조건을 만족하는 Key가 없는 경우 -1을 출력한다. 만약 해당하는 Key가 두 개 이상 존재한다면 ?를 출력한다. 모든 출력은 개행을 포함해야 한다.

출력

각 줄에 걸쳐 3번 명령에 대한 결과를 출력한다.

코드

from bisect import insort
import sys
input = sys.stdin.readline

n, m, k = map(int, input().split())

jbnu = dict()
keys = []
# 초기 데이터 입력
for _ in range(n):
    key, value = map(int, input().split())
    jbnu[key] = value
    keys.append(key)

keys.sort()


def search(target):  # 검색
    s, e = 0, len(keys)-1
    gap = sys.maxsize
    res = [-1]
    while s <= e:
        mid = (s+e)//2
        if jbnu[keys[mid]] == target:
            return [target]
        else:
            mid_key = keys[mid]
            if abs(mid_key-target) <= k:
                if abs(mid_key-target) < gap:
                    gap = abs(mid_key-target)
                    res = [mid_key]
                elif abs(mid_key-target) == gap:
                    res.append(mid_key)
            if mid_key < target:
                s = mid+1
            else:
                e = mid-1
    return res


# 명령 입력
for _ in range(m):
    cmd = list(map(int, input().split()))
    res = search(cmd[1])
    if cmd[0] == 1:  # 데이터 추가
        if res[0] != -1 and jbnu[res[0]] == cmd[1]:
            continue
        else:
            insort(keys, cmd[1])
            jbnu[cmd[1]] = cmd[2]
    elif cmd[0] == 2:   # 데이터 변경
        res = search(cmd[1])
        if res[0] >= 0 and len(res) == 1:
            jbnu[res[0]] = cmd[2]
    else:   # 데이터 출력
        res = search(cmd[1])
        if len(res) > 1:
            print("?")
        else:
            if res[0] >= 0:
                print(jbnu[res[0]])
            else:
                print(-1)

 

• search()
  데이터를 검색할 때는 이분탐색을 이용하여 검색하려는 key를 찾고 만약 없다면 거리가 K 이하인 다른 key들 중 가장 가까운 키를 리턴한다. 이때 가까운 키는 2개 이상일 수도 있다.

이분탐색으로 데이터를 찾는다는 풀이 방법을 찾는 것은 어렵지 않았지만 시간 초과 때문에 꽤 오래 걸렸다.

아마 데이터 추가 부분에서 처음에는 단순히 새로운 데이터를 추가하고 다시 정렬하는 방식으로 했기 때문이었던 것 같다.

그래서 데이터를 추가할 때도 이분탐색을 이용하도록 했다. bisect 라이브러리의 insort 메소드는 1번째 인자로 주어진 배열을 이분탐색하여 2번째 인자가 들어갈 위치를 찾아서 넣어준다.

이렇게 데이터 추가 방법을 바꿨더니 시간 초과를 해결할 수 있었다.