[Python] 백준 14428 - 수열과 쿼리 16

https://www.acmicpc.net/problem/14428

문제

길이가 N인 수열 A1, A2, ..., AN이 주어진다. 이때, 다음 쿼리를 수행하는 프로그램을 작성하시오.

• 1 i v : Ai를 v로 바꾼다. (1 ≤ i ≤ N, 1 ≤ v ≤ ((10^9)))
• 2 i j : Ai, Ai+1, ..., Aj에서 크기가 가장 작은 값의 인덱스를 출력한다. 그러한 값이 여러개인 경우에는 인덱스가 작은 것을 출력한다. (1 ≤ i ≤ j ≤ N, 1 ≤ v ≤ ((10^9)))

수열의 인덱스는 1부터 시작한다.

입력

첫째 줄에 수열의 크기 N이 주어진다. (1 ≤ N ≤ 100,000)

둘째 줄에는 A1, A2, ..., AN이 주어진다. (1 ≤ ((A_i)) ≤ (10^9))

셋째 줄에는 쿼리의 개수 M이 주어진다. (1 ≤ M ≤ 100,000)

넷째 줄부터 M개의 줄에는 쿼리가 주어진다.

출력

2번 쿼리에 대해서 정답을 한 줄에 하나씩 순서대로 출력한다.

 

---

풀이

세그먼트 트리의 기본 응용문제이다.

본격적으로 트리의 노드에 정수 말고 다른 값들을 저장하기 시작한다.

 

블로그에 Bottom-Up 구현도 정리했겠다, 이 문제는 재귀 안쓰고 풀어보려고 한다.

 

사실 재귀를 안쓰면 update 함수의 구현이 살짝 방법이 달라진다. 왜냐하면, 예를들어 기존에 최솟값이었던 값이 삭제되고 그 값이 큰 값으로 바뀔 경우 기존의 방법으로는 이를 해결해줄 수가 없다.

그래서 살짝 나이브하게, 맨 처음에 리프 노드의 값을 바꿔주고 인덱스를 2로 나눠서 상위 노드로 올라가준 뒤

인덱스가 1 이상이면 자식 노드 두개(idx*2, idx*2+1)을 직접 비교해서 최솟값을 지정해준다. 이렇게 해야 정확하게 데이터가 업데이트 된다.

 

update 함수의 구현은 그러면 다음과 같이 할 수 있을 것이다.

 

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

 

문제는 search 함수이다.

최솟값을 출력하는 것이 아니라, 최솟값의 인덱스를 출력한다. 그것도 최대한 왼쪽에 있는 거 우선으로.

 

사실 이 문제의 해결법은 매우매우 간단하다.

세그먼트 트리의 각 노드에 최솟값만을 저장하는게 아니라, 튜플 형태로 (최솟값, 인덱스)를 같이 저장하는 것이다. 튜플의 비교 방식에서 첫 번째 원소가 같다면 자동으로 두 번째 원소를 비교하기 때문에 최솟값 연산을 해준다면 늘 인덱스가 가장 작은 놈이 저장됨을 알 수 있다.

 

만약 아무것도 없는 노드라면, (9999999999, -1)를 기본값으로 가지고 있게 하도록 하자.

 

그러면 전체 코드는 다음과 같다.

 

import sys
input = sys.stdin.readline

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

size = 1
while size < N:
    size <<= 1
tree = [(9999999999, -1)]*(2*size)
size -= 1

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

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

def search(l, r):
    l += size
    r += size
    ans = (9999999999, -1)
    while l <= r:
        if l % 2 == 1:
            ans = min(ans, tree[l])
            l += 1
        l //= 2

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

M = int(input())
for i in range(M):
    o, i, j = map(int, input().split())
    if o == 1:
        update(i, (j, i))
    else:
        print(search(i, j))

 

이 문제는 배열의 인덱스가 1부터 시작하므로, 트리를 맨 처음 만들 때 idx -= 1 을 해주었다. 주의하도록 하자.

 

1트컷