본문 바로가기

코딩테스트

[코드트리] 삼성 코테 2024 하반기 - 코드트리 DB

구현 코드

import bisect

q = int(input())
queries = []
for _ in range(q):
    each = input()
    queries.append(each)

name_to_value = {}
value_to_name = {}
sorted_values = []

for query in queries:
    if query == 'init':
        name_to_value = {}
        value_to_name = {}
        sorted_values = []

    elif 'insert' in query:
        _, name, value = query.split()
        value = int(value)

        if name in name_to_value or value in value_to_name:
            print(0)
        else:
            name_to_value[name] = value
            value_to_name[value] = name
            bisect.insort(sorted_values, value)  # 정렬 유지
            print(1)

    elif 'delete' in query:
        _, name = query.split()

        if name in name_to_value:
            value = name_to_value.pop(name)
            value_to_name.pop(value)

            idx = bisect.bisect_left(sorted_values, value)
            del sorted_values[idx]
            print(value)
        else:
            print(0)

    elif 'rank' in query:
        _, k = query.split()
        k = int(k)

        if len(sorted_values) < k:
            print("None")
        else:
            kth_value = sorted_values[k - 1]
            print(value_to_name[kth_value])

    elif 'sum' in query:
        _, k = query.split()
        k = int(k)

        idx = bisect.bisect_right(sorted_values, k)
        result = sum(sorted_values[:idx])
        print(result)

bisect

시간 복잡도를 줄이기 위해 python의 bisect 라이브러리를 활용하였다.

bisect 모듈이란 정렬된 배열에서 이진 탐색을 쉽게 구현할 수 있도록 도와주는 모듈이다.

기존 코드에서 insert, delete, rank, sum에서 해당 요소에 접근하느라 배열을 모두 순회하여 O(n) 시간복잡도를 가졌었다.

bisect 주요 함수

  • bisect.bisect_left(a,x) : x가 들어갈 수 있는 가장 왼쪽 인덱스 반환
  • bisect.bisect_right(a,x) : x가 들어갈 수 있는 가장 오른쪽 인덱스 반환
  • bisect.insort(a,x) : a 리스트에 x를 정렬 유지하며 삽입

개선 방법

  • 값을 삽입할 때는 bisect.insort 로 정렬 유지
  • 값을 삭제할 때는 bisect.bisect_left 로 위치 찾아서 삭제
  • k번째 값을 찾을 때는 정렬된 리스트에서 인덱스 접근
  • 합계를 구할 때는 bisect 로 범위 찾아서 slicing sum