백준 14897번: 서로 다른 수와 쿼리 (Segment Tree)

문제링크: 14897번: 서로 다른 수와 쿼리 1 (acmicpc.net)

14897번: 서로 다른 수와 쿼리 1

이 문제는 배열 A의 길이 N이 N <= 1,000,000 일때, 배열의 각 숫자의 범위가 <=1,000,000,000 일때, 쿼리의 개수 Q가 Q <= 1,000,000 일때 (l, r)쿼리가 주어졌을 때 l부터 r까지의 범위 내 서로 다른 숫자가 몇개인지 반환하는 쿼리를 만드는 문제이다. 일단 N과 Q가 1,000,000 으로 꽤 큰편에 속해 최적화를 잘 해줘야 한다. 일단 배열의 각 숫자가 최대 1,000,000,000 인데 반해 배열의 길이는 1,000,000이므로 좌표 압축을 하여 각 숫자를 비교하는데 걸리는 시간을 줄일 수 있다.

 

처음에는 비슷한 문제를 mo's 알고리즘을 사용해 풀어 보았기에 sqrt decomposition 을 사용해 쿼리를 정렬하고, mo's 알고리즘을 사용했지만 시간 초과가 나왔다. mo's 알고리즘의 시간 복잡도를 계산하면 O( (Q+N)√N)으로 2,000,000 * 1,000 으로 20억번의 연산을 해야 하므로 시간 내에 풀 수 없을 것이다.

 

고민하던 중 질문 게시판에서 힌트를 받아 해당 인덱스의 숫자와 같은 숫자가 나타나는 최소 다음 인덱스를 저장하는 배열 nxt를 만들어서 풀 수 있다는 것을 알았다. 예시로 (1,2,1,3,1)이 있다면 nxt는 (2,-,4,-,-)로 표현할 수 있다. 2, 3, 마지막 1은 배열에 없으므로 -로 표현했지만 계산에서는 인덱스의 범위인 [0, N-1] 보다 큰 N을 할당한다. 그렇게 nxt를 만들면 (2,5,4,5,5) 로 될 것이다.

 

l과 r을 인덱스라고 쳤을때 (쿼리의 l, r -> l-1, r-1로 인덱스로 변환), 쿼리를 수행할 때 l부터 r까지의 nxt를 보면서 nxt의 값이 r보다 큰 값들의 갯수를 세어주면 된다. r보다 작은 값들은 l ~ r 범위 내에서 한 번 더 나온다는 뜻이므로 카운트 해주지 않고, r보다 큰 값들은 해당 숫자와 중복되는 숫자의 다음 인덱스가 r 범위 뒤에서 나온다는 뜻이므로 해당 범위 내에서는 마지막 중복이 될 것이다. 예시로 아까 (2,5,4,5,5)에서 l:0 r:3를 본다면 (2,5,4,5)에서 r보다 큰 nxt: 5, 4, 5 -> A: 2, 1, 3로 답은 3이 된다. 

 

그러면 이제 지난 수열과 쿼리 1문제와 비슷하게 머지 소트 트리를 사용하여 풀어 보았지만 시간 초과를 받았다. 시간 복잡도를 계산해 보면 일단 머지 소트 과정 O(N logN), M개의 쿼리에서 머지 소트 리스트를 찾는 과정 + 이분 탐색 과정 O(M * logN * log N) 이 되어 O(M log^2N)이 된다. 두개를 계산하면 O((N + M logN) logN) 로 log(백만) 대략 20이므로 (21,000,000) * 20으로 4억번의 연산으로 아까보다는 적어졌지만 아직 충분하지 않은가보다.

 

 그래서 다시 사용한 방법이 오프라인 쿼리를 사용하는 것이다. 오프라인 쿼리를 사용하는 방법은 먼저 배열을 그대로 쓰지 않고 bool배열 (0,0,0,0,0..)에서 각 인덱스의 숫자를 추가할 때마다 1로 처리하는 쿼리를 만들고, 개수를 세는 것을 구간합을 구하는 쿼리로 만드는 것이다. 여기서 추가 쿼리는 (인덱스, 값)으로 저장하고 구간합 쿼리는(l, r, 쿼리순서)로 저장한다. 그리고 두개의 쿼리를 추가 쿼리는 값, 구간합 쿼리는 r을 기준으로 내림차순으로 정렬하고, 값과 r이 같다면 우리는 r보다 큰 값을 세야 하므로 구간합 쿼리를 우선적으로 정렬한다. 정렬된 쿼리를 오프라인으로 수행하면 답을 얻을 수 있다.

import sys

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

nums = list(map(int,sys.stdin.readline().rstrip().split()))
nxt = [N for _ in range(N)] #해당 인덱스의 숫자가 나오는 다음 인덱스

nxtcnt = {}
for idx, val in enumerate(nums):
    if val in nxtcnt.keys():
        nxt[nxtcnt[val]] = idx
        nxtcnt[val] = idx
    else:
        nxtcnt[val] = idx
tree = [0 for _ in range(4*N)]

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

def return_sum(start,end,index,i,j):
    if i <= start and end <= j:
        return tree[index]

    if start > j or end < i:
        return 0

    mid = (start+end)//2

    if i > mid:
        return return_sum(mid+1, end, 2*index+1, i, j)
    elif j <= mid:
        return return_sum(start, mid, 2*index, i, j)

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

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

query = []

for idx, val in enumerate(nxt):
    query.append((0,idx,val)) #segment tree에 값 할당 맨앞 0:구간합 쿼리 우선

for i in range(M):
    a,b = map(int,sys.stdin.readline().rstrip().split())
    query.append((1,a-1,b-1,i)) #segment tree 구간합 맨앞 1:구간합 쿼리 우선

anslst = [0 for _ in range(M)]
query.sort(key = lambda x: (x[2], x[0]), reverse = True)

for x in query:
    if x[0] == 0:
        update_tree(0, N-1, 1, x[1])
    else:
        anslst[x[3]] = return_sum(0, N-1, 1, x[1], x[2])

for x in anslst:
    print(x)

풀이에서는 update_tree 쿼리는 맨앞에 0을 주고 return_sum 쿼리는 맨앞에 1을 주어 두 쿼리간의 우선순위를 비교를 했다. 예시로 아까 nxt(2,5,4,5,5) 가 있다고 치고 쿼리가 l:0 r:3이라면 . 그러면 추가 쿼리

update_tree()는 (0,0,2), (0,1,5), (0,2,4), (0,3,5), (0,4,5) 가 생길 것이고 구간합 쿼리

return_sum()는 (1,0,3,0) 한개 존재하게 될 것이다. 

쿼리 순서는:

(0,4,5) - (0,3,5) - (0,1,5) - (0,2,4) - (1,0,3,0) - (0,0,2) 순으로 진행된다.

그러면 처음 bool 배열 (0,0,0,0,0) 에서

(0,0,0,0,1) -> (0,0,0,1,1) -> (0,1,0,1,1) -> (0,1,1,1,1) -> 구간합 (l:0~r:3) = 3 -> (1,1,1,1,1)

이렇게 진행되어 구간합이 서로 다른 수를 잘 세는 것을 볼 수 있다.

 

위 과정은 시간 복잡도가 쿼리의 개수가 (N+M)개로 늘어났지만, return sum 쿼리와 update 쿼리가 모두 logN 선에서 해결 가능하므로 O((N+M) log (N)) ~= 40,000,000 으로 아까보다 더 빠른 시간 복잡도를 가진다. 근데 그런데도 처음에 시간 초과가 나서 update_tree()함수를 재귀 대신 while문을 사용해 최적화 했더니 맞았습니다를 받을 수 있었다.