백준 13537번: 수열과 쿼리 1 (Merge sort tree)

문제링크:13537번: 수열과 쿼리 1 (acmicpc.net)

13537번: 수열과 쿼리 1

이 문제는 길이가 N인 수열이 있을 때, 쿼리의 입력이 (i, j, k) 이렇게 주어지면 i번째부터 j번째 까지의 수열의 원소 중 k 보다 큰 값들의 개수를 출력하는 문제이다. 여기서 N의 최대값은 100,000, 쿼리의 개수의 최대값도 100,000 이기 때문에 불러오는 쿼리의 시간 복잡도를 O(logN) 선에서 처리해야 한다. 

 

나는 전에 몇번 다루었던 merge sort를 활용하고 이분 탐색을 활용하여 문제를 풀었다. merge sort는 세그먼트 트리 형태로 범위를 절반씩 분할 해 가며 sort한다. 예시로 [8,7,6,5,4,3,2,1] 의 8개의 원소를 가진 배열을 merge sort한다고 치면 left 자식에는 인덱스가 0~4 까지, right는 5~8 까지 나누어 간다. left는 다시 (L): 0~2, (R): 3~4 로 나누어 진다. 여기서 (L)을 다시 나누면 (L): 0, (R): 1의 인덱스 범위를 갖게 된다.

 

여기서 merge해가며 정렬된 리스트를 만들어 가는데, 나는 이 문제를 위해서 각 노드에서 merge한 리스트들을 tree에 저장하였다. 이 리스트는 각 노드의 인덱스 범위의 정렬된 리스트라는 의미를 갖는다.

 

우리가 찾는 쿼리 i, j의 범위가 노드의 범위와 일치한다면, 그 노드의 리스트에서 k보다 큰 수를 찾으면 되는 것이다. 만약 범위가 다르다면, 기존의 세그먼트 트리 방식을 이용하여 일치하는 범위의 노드들에서 k보다 큰 값들을 더해가며 i 부터 j까지 k보다 큰 값의 개수의 총합을 구할 수 있다. 여기서 리스트에서 k보다 큰 수를 찾는 방법은 이분 탐색을 이용하였다.

 

근데 세그먼트 트리 자료구조에서 범위에 맞는 노드를 발견하고, 거기서 다시 이분 탐색을 하면 시간 복잡도가 O(log^2(N)) 이 되는 것이 아닌가? 라고 생각 할 수 있다. 그렇지만 생각해 보면 범위에 맞는 노드를 찾기 위해 트리를 깊게 들어갈 수록, 해당 노드의 리스트 길이는 절반씩 줄어들게 된다.

 

만약 N=100,000 이고 (i, j)가 (1, 100,000)이라고 치면 범위를 포함하는 노드는 트리의 루트가 된다. 그러므로 여기서 해당 노드를 찾는데 시간 복잡도 O(1), + 이분 탐색하는데 걸리는 시간 O(log(100,000)) 이 된다. 반대로 만약 (i, j)가 (1,1)이라면 트리의 가장 깊은 노드까지 가는데 시간 O(log(100,000)) + 이분 탐색 O(1)이 된다.

 

그럼 중간 중간으로하면 최대값이 나오지 않을까? N=100,000일때 트리의 깊이는 17이 되어야 한다. (log2(100,000) = 16.6 ) 중간 깊이까지 9만큼 간다고 하면 리스트의 길이는 100,000 // (2^9) = 195 일것이고 log2(195) = 7.67 이므로 탐색한 시간은 9+8 으로 17 (O(log(100,000))) 가 된다.

 

내가 생각하는 최악의 경우는 (i,j) 가 (50000,500001) 일때이다. 이때는 범위 (50000, 50000), (50001,50001) 까지 세그먼트 트리에서 최대 깊이까지 2번 가야 하기 떄문에, 시간 복잡도가 2 * (O(logN) * O(1) )이 될 것이다.

 

함수를 어떻게 만드냐에 따라 실제 걸리는 시간이 완벽하게 모두 일치하지는 않겠지만, 결국 범위를 어떻게 잡던 간에 쿼리의 시간 복잡도는 항상 O(logN) 선에서 처리가 가능하게 된다.

import sys

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

nums=list(map(int,sys.stdin.readline().rstrip().split()))

tree=[[] for _ in range(4*N)]

def merge_sort(start,end,index):
    if start==end:
        tree[index] = [nums[start]]
        return tree[index]

    mid = (start+end)//2

    tree[index] = merge(merge_sort(start,mid,2*index), merge_sort(mid+1,end,2*index+1))
    return tree[index]

def merge(left,right):
    a=0
    b=0
    sortlst=[]
    
    while a<len(left) and b<len(right):
        if left[a]<=right[b]:
            sortlst.append(left[a])
            a+=1
        else:
            sortlst.append(right[b])
            b+=1

    if a==len(left):
        for x in right[b:]:
            sortlst.append(x)
    else:
        for x in left[a:]:
            sortlst.append(x)
            
    return sortlst

def binary_search(series,start,end,val):
    if start==end:
        if series[start]<=val:
            return start+1
        return start
    
    mid = (start+end)//2
    
    if series[mid] > val:
        return binary_search(series,start,mid,val)
    else:
        return binary_search(series,mid+1,end,val)
    
def find_ans(start,end,index,i,j,k):

    if i <= start and end <= j:
        idx=binary_search(tree[index],0,len(tree[index])-1,k)
        return len(tree[index]) - idx

    if end < i or j < start:
        return 0

    mid = (start+end)//2

    return find_ans(start,mid,2*index,i,j,k) + find_ans(mid+1,end,2*index+1,i,j,k)

merge_sort(0,N-1,1)

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

for _ in range(M):
    a,b,c=map(int,sys.stdin.readline().rstrip().split())
    print(find_ans(0,N-1,1,a-1,b-1,c))

세그먼트 트리를 이용하여 N개의 자료에 대해 각종 연산들을 취하거나 반환하는 쿼리들을 시간 복잡도 O(logN) 안에 처리할 수 있다는 것이 놀랍고, 풀떄마다 새로워서 재미있는 문제였다.

 

여담으로 이분 탐색 binary search()함수를 이용해 해당 위치의 sortlst에서 k보다 큰 값이 처음으로 나오는 인덱스 (upper bound) 를 찾는 함수를 만들기 위해 이것저것 해보다가 문제의 채점 오류를 발견했다. 주어지는 수열에 중복 원소가 있고, 그 중복 원소가 쿼리의 k일때 올바르지 못한 답을 출력해도 정답 처리가 된다는 것을 알았다. 요청 게시판에 올려 데이터를 추가할 수 있었다. 처음으로 내가 데이터를 추가해 본 문제였다!