백준 12899번: 데이터 구조 (Segment Tree)

문제링크:12899번: 데이터 구조 (acmicpc.net)

12899번: 데이터 구조

이 문제는 S라는 데이터베이스에 1<=N<=2,000,000 인 자연수를 저장하고, K번째로 작은 수를 반환하고 제거할 수 있는 쿼리를 만드는 문제이다.

 

처음 문제를 봤을때는 2,000,000개의 케이스를 처리하기 위해 무조건 시간 복잡도가 O(logN) 인 쿼리를 만들어야 하기에 세그먼트 트리나 이분 탐색등을 이용해야 할 것이라고 생각했지만, 어떻게 이용할 지 감이 잘 안잡혔다. 그러던 중 자연수의 범위가 2,000,000 밖에 되지 않아 200만개의 자리만큼의 세그먼트 트리를 만들고 추가하고, 탐색하는 쿼리만 만들면 될 것이라고 생각했다.

 

내가 생각한 방법은 먼저 0부터 2,000,000 까지의 리스트를 만들고 각 자연수가 하나 추가될때마다 +=1, 빠질 때마다 -=1 해주었다. 그리고 세그먼트 트리는 트리의 각 start, end 범위가 그 범위 내의 자연수 개수를 저장하도록 하였다. K번째 작은 수를 탐색할 때에는 내가 탐색하고자 하는 K번째 수를 탐색할 때 left가 K보다 크다면 K번쨰 수는 left범위에 있다는 뜻이므로 left쪽을 탐색하고, left가 K보다 작다면 right 범위에 있다는 뜻이므로 right로 탐색하고,  left의 자연수 개수만큼 스킵했으니 K에서 left만큼 빼주고 탐색을 이어나간다.

 

처음에는 재귀함수를 이용해 숫자를 insert, K번째 작은 숫자를 pop 하는 쿼리를 만들었는데 6400ms로 생각보다 시간이 많이 느려서 while 문을 이용해서도 만들었더니 시간이 약 1/4정도로 크게 줄었다. 또 처음엔 0부터 2,000,000까지의 리스트를 만들고 트리를 또 만들었는데 생각해 보니 현재 데이터베이스의 모든 값들의 상태를 보는것이 아니라면, 리스트를 굳이 만들 필요가 없어서 tree[] 리스트 내에서 해결했다.

 

tree는 N의 최대 정수 범위가 2,000,000이므로 2,000,000보다 큰 최소 2의 제곱수 2^(21) 의 두배인 2^(22)로 정해주면 2,000,000을 커버할 수 있는 최소의 완전 이진 트리 범위로 만들 수 있다.

 

이용한 방법은 insert()함수에서는 먼저 해당 value까지 접근 할 때 모든 인덱스들은 범위 안에 value를 포함하므로 접근하면서 마주치는 모든 tree[index]에 +=1 해준다. 그리고 start==end 할 때 start, end, value가 다 일치하는 마지막 탐색 범위이므로 탐색을 끝내준다.

 

반대로 pop_order()함수에서는 해당 value까지 접근 할 때 마주치는 모든 tree[index]에 -=1 해준다. start==end 할때에 마찬가지로 해당 value의 마지막 탐색 범위이므로 해당 숫자를 return 해주어 탐색을 끝내준다.

 

재귀를 이용한 코드:

import sys

N=int(sys.stdin.readline().rstrip())

tree=[0 for _ in range(1<<22)]

def insert(start,end,index,val):
    if start==end:
        tree[index] += 1
        return tree[index]
    
    mid = (start+end)//2

    if val<=mid:
        tree[index] = insert(start,mid,2*index,val) + tree[2*index+1]
    else:
        tree[index] = tree[2*index] + insert(mid+1,end,2*index+1,val)
    return tree[index]

def pop_order(start,end,index,order):
    
    tree[index]-=1
    
    if start==end:
        return start
    
    left = tree[2*index]
    right = tree[2*index+1]
    mid = (start+end)//2
    
    if order <= left:
        return pop_order(start,mid,2*index,order)
    else:
        return pop_order(mid+1,end,2*index+1,order-left)
    
for _ in range(N):
    a,b=map(int,sys.stdin.readline().rstrip().split())

    if a==1:
        insert(1,2_000_000,1,b)
    else:
        print(pop_order(1,2_000_000,1,b))

 

while 을 이용한 코드:

import sys

N=int(sys.stdin.readline().rstrip())

tree=[0 for _ in range(1<<22)]

def insert(start,end,index,val):
    while 1:
        tree[index]+=1
        
        if start==end:
            break
            
        mid = (start+end)//2
        
        if val<=mid:
            end = mid
            index*=2
        else:
            start = mid+1
            index= 2*index + 1
    return

def pop_order(start,end,index,order):

    while 1:
        tree[index]-=1

        if start==end:
            return start
        
        left = tree[2*index]
        right = tree[2*index+1]
        mid = (start+end)//2
        
        if order <= left:
            end = mid
            index*=2
        else:
            start = mid+1
            index = 2*index + 1
            order-=left
    
for _ in range(N):
    a,b=map(int,sys.stdin.readline().rstrip().split())

    if a==1:
        insert(1,2_000_000,1,b)
    else:
        print(pop_order(1,2_000_000,1,b))

재귀를 이용하면 코드를 읽고 짜는데 더 가독성이 좋지만 시간이 느려진다는 단점이 있고, 범위가 많아지면 sys.recursionlimit() 를 통해 재귀 깊이를 정해줘야 한다. 재귀 깊이를 정할 때 또 너무 많이 정하면, 정한 그 깊이에 비례하여 더 많은 메모리을 소모한다는 단점이 있다. 이럴 땐 while을 통해서 더 효율적으로 코드를 짤 수 있다. 가독성은 별로 안좋지만 while을 사용한 코드가 더 깔끔한 것 같다.