[백준 14428] 수열과 쿼리 16

문제 출처 : https://www.acmicpc.net/problem/14428

14428번: 수열과 쿼리 16

세그먼트 트리를 조금 응용한 문제 입니다. 기존에는 리스트의 값을 출력 하였다면 이 문제에서는 해당 인덱스를 출력하는 문제 입니다. 이처럼 세그먼트 트리를 어떤 값을 찾는 것이 아니라 보조적인 수단으로 사용하는 문제가 많이 출제되니 이런 케이스의 문제를 많이 풀어봐야 합니다.

이 문제를 푸는 방법은 기존의 세그먼트 트리와 많이 유사합니다. 리스트의 인덱스를 저장할 새로운 트리를 만들어도 되고, 처음부터 리스트에 인덱스를 저장하는 방법도 있습니다. 또는 바텀업 세그먼트 트리를 사용하여 찾아도 됩니다. 자신이 가장 익숙한 방식으로 진행하면 됩니다. 여기서는 바텀업 세그먼트 트리를 사용하여 문제를 해결해 보겠습니다. 일반적인 세그먼트 트리로 연습을 충분히 한 후 바텀업 세그먼트 트리를 연습해 보시기 바랍니다.

입력 받기

mii = lambda : map(int, input().split())
N = int(input())
arr = list(mii())

M = int(input())

for _ in range(M):
    a, b, c = mii()    
    if a == 1:        
        pass
    elif a == 2:        
        pass

수열의 크기 N, 수열 arr, 쿼리의 개수 M, M개의 쿼리를 입력 받습니다. 지금은 쿼리 부분은 pass로 넘어가도록 하였습니다. 함수를 만들고 나머지 부분을 채워놓도록 하겠습니다.

트리 만들기

size = 1
while size < N:
    size <<= 1

INF = float("INF")
tree = [[INF, INF] for _ in range(2 * size)]

N의 크기에 맞게 트리를 만들어 줍니다. 반복문을 통해 말단 노드에 들어가는 크기를 만들어 줍니다. 트리의 크기를 계산을 통해 알고 싶다면 다빈치코딩 알고리즘의 세그먼트 트리의 크기 계산을 확인해 보시기 바랍니다.

업데이트 함수 만들기

def update(i, num):
    idx = i + size
    tree[idx] = [num, i]
    idx //= 2
    while idx:
        tree[idx] = min(tree[2 * idx], tree[2 * idx + 1])
        idx //= 2

for i in range(N):
    update(i + 1, arr[i])

바텀업 세그먼트 트리에는 init 함수가 따로 없습니다. 반복문을 통해 모든 값을 업데이트 하는 것이 초기화 입니다.

idx = i + size

update함수에 들어가는 i는 리스트의 인덱스 번호 입니다. i에 size를 더해주어 말단에 위치하도록 하였습니다.

tree[idx] = [num, i]

다음으로 tree에는 값인 num과 인덱스 번호인 i를 넣어줍니다.

idx //= 2

이제 트리의 위로 올라가기 위해 인덱스 값을 반으로 줄여 줍니다.

while idx:

다음으로 while 반복문을 통해 트리를 구성합니다.

tree[idx] = min(tree[2 * idx], tree[2 * idx + 1])

min 함수로 최소값을 저장 합니다. 트리에는 비교할 값과 인덱스가 저장되어 있습니다. min 함수로 비교할 때 key가 없으면 앞의 값을 비교합니다.

idx //= 2

인덱스 값을 반으로 줄여가며 최상단까지 만들어 줍니다.

for i in range(N):
    update(i + 1, arr[i])

init 함수 대신 반복문으로 전체 트리를 구성합니다. 문제에서 리스트의 인덱스는 1부터 시작하기 때문에 for문의 업데이트 함수에서 i + 1로 인덱스를 1부터 시작하도록 하였습니다.

최소값 인덱스 구하는 함수 만들기

다음으로 최소값의 인덱스를 리턴하는 query 함수를 만들겠습니다.

def query(left, right):
    left += size
    right += size 
    ans = [INF, INF]
    
    while left <= right:
        if left % 2 == 1:
            ans = min(ans, tree[left])
            left += 1
        left //= 2

        if right % 2 == 0:
            ans = min(ans, tree[right])
            right -= 1
        right //= 2
    
    return ans[1]

query 함수를 설명하겠습니다.

  	left += size
    right += size 
    ans = [INF, INF]

먼저 left와 right 인덱스에 size를 더해 말단 노드에 위치하도록 합니다. 다음으로 ans의 초기값을 [INF, INF]로 하여 최대값으로 지정해 주었습니다.

    while left <= right:
        if left % 2 == 1:
            ans = min(ans, tree[left])
            left += 1
        left //= 2

반복문은 left와 right가 역전하기 전까지 반복을 진행 합니다. left는 홀수일 경우 해당 값을 확인합니다. 이해가 잘 안간다면 다빈치코딩 알고리즘의 바텀업 세그먼트 트리 부분을 다시 확인해 보기 바랍니다.

        if right % 2 == 0:
            ans = min(ans, tree[right])
            right -= 1
        right //= 2

    return ans[1]

right는 짝수일 경우 해당 값이 최소값인지 확인합니다. 반복문이 완료 되면 ans값의 값이 아닌 인덱스가 저장되어 있는 ans[1]을 리턴해 줍니다.

쿼리문 만들기

for _ in range(M):
    a, b, c = mii()    
    if a == 1:        
        update(b, c)
    elif a == 2:        
        print(query(b, c))

마지막으로 a, b, c값에 따라 a가 1인 경우는 업데이트를, 2인 경우에는 구간의 최소값의 인덱스를 출력하도록 합니다. 인덱스의 시작을 1이 되도록 하였기 때문에 b나 c의 값에 1을 빼지 않았습니다.

전체 코드

전체 코드를 확인해 보겠습니다.

mii = lambda : map(int, input().split())
N = int(input())
arr = list(mii())

M = int(input())
INF = float("INF")

size = 1
while size < N:
    size <<= 1

tree = [[INF, INF] for _ in range(2 * size)]

def update(i, num):
    idx = i + size
    tree[idx] = [num, i]
    idx //= 2
    while idx:
        tree[idx] = min(tree[2 * idx], tree[2 * idx + 1])
        idx //= 2

for i in range(N):
    update(i + 1, arr[i])

def query(left, right):
    left += size
    right += size 
    ans = [INF, INF]
    
    while left <= right:
        if left % 2 == 1:
            ans = min(ans, tree[left])
            left += 1
        left //= 2

        if right % 2 == 0:
            ans = min(ans, tree[right])
            right -= 1
        right //= 2
    
    return ans[1]

for _ in range(M):
    a, b, c = mii()    
    if a == 1:        
        update(b, c)
    elif a == 2:        
        print(query(b, c))