백준 14427 수열과 쿼리 15 (파이썬)

문제 : https://www.acmicpc.net/problem/14427

 

 

제출코드 (정답)

import sys
input = sys.stdin.readline
print = sys.stdout.write

class Node():
    def __init__(self, index, value):
        self.index = index
        self.value = value

def update(index, value):
    global base, N, arr, M, base, cnt, segmentTree
    p = base + index
    segmentTree[p] = Node(index, value)
    p //= 2
    while p >= 1:
        left, right = p * 2, p * 2 + 1
        # 최소값 우선 순위
        # 수열에서 크기가 가장 작은 값의 인덱스를 출력 그러한 값이 여러 개인 경우에는 인덱스가 작은 것을 출력한다.
        if segmentTree[left].value == segmentTree[right].value:
            if segmentTree[left].index < segmentTree[right].index:
                segmentTree[p] = segmentTree[left]
            else: segmentTree[p] = segmentTree[right]

        elif segmentTree[left].value < segmentTree[right].value:
            segmentTree[p] = segmentTree[left]

        else: segmentTree[p] = segmentTree[right]
        p //= 2

def insert(index, value):
    update(index, value)

def init():
    global base, N, arr, M, base, cnt, segmentTree

    N = int(input())  # 수열의 크기
    arr = [0] + list(map(int, input().split()))
    M = int(input())  # 쿼리의 개수
    cnt = base = 1  # 부모 노드의 개수들
    segmentTree = []
    MAX = int(10e9)

    # N이 6이면 base = 2^h-1 = 2^3-1 = 7
    while base < N:
        base *= 2
        cnt += base  # 세그멘트 트리의 노드 개수
    base -= 1

    # 세그멘트 트리 초기화, 배열로 구현
    for i in range(cnt+1):
        segmentTree.append(Node(i, MAX))

if __name__ == "__main__":
    global base, N, arr, M, base, cnt, segmentTree
    init()

    # 리프노드에 수열 넣기
    for i in range(1, N + 1): insert(i, arr[i])

    for i in range(M):
        query = list(map(int, input().split()))
        if query[0] == 2:  # 수열에서 크기가 가장 작은 값의 인덱스를 출력 그러한 값이 여러개인 경우에는 인덱스가 작은 것을 출력한다.
            print(f'{segmentTree[1].index}\n')
        else:
            update(query[1], query[2])  # Ai를 v로 바꾼다

 

반응형

풀이

세그멘트 트리를 처음 배워서 이를 적용하여 문제를 풀어 보았습니다.

 

Node 클래스에는 index와 value를 변수로 두어 인덱스와 값을 저장하게 만들었습니다.

저는 init, insert, update 함수들을 구현하여 문제를 해결했습니다.  

init : 세그멘트 초기화를 담당

insert : 초기 리프 노드에 수열과 인덱스 삽입

update : 리프 노드의 변경이 이루어 질 때마다 관련 부모 노드들을 업데이트

 

예를 들어, 예제 1의 수열 5 4 3 2 1을 세그멘트 트리에 넣는 init를 실행하면 다음처럼 세그멘트 트리가 만들어집니다.

 

먼저 완전 이진 트리를 만들고 왼쪽부터 리프 노드에 수열을 넣습니다.

그 후 부모 노드는 value가 작은 node가 됩니다. 

 

세그멘트의 개념은 트리로 설명을 하지만 구현은 배열로 합니다.

이는 완전 이진 트리라서 부모 트리에서 *2 를 하면 왼쪽 자식 노드, *2+1을 하면 오른쪽 자식 노드가 바로 나오고 크기가 딱 정해지기 때문에 배열의 범위를 절대 벗어나지 않기 때문에 배열로 구현이 가능합니다.

 

update 와 insert 함수는 동일한데 굳이 insert 함수를 만든 이유는 코드의 가독성이 높아지고 다음에 더 어려운 문제가 나오면 위 함수들을 적절히 변경해서 문제를 해결할 수 있게 해주기 때문입니다.

 

쿼리의 첫번째 값이 1이면 Ai를 v로 값을 바꿔주고 이와 관련된 모든 부모 노드의 정보를 업데이트 해줍니다.

쿼리의 첫번째 값이 2이면 세그멘트 트리의 루트 노드의 index를 출력해줍니다. 

 

결과 및 정리

다른사람의 풀이를 보니 heapq를 써서 쉽게 풀었고 위 문제는 우선순위 큐를 사용하면 쉽게 풀릴거라고 생각했지만 저는 세그멘트 트리를 배웠어서 이를 사용해서 꼭 풀어보고 싶어 우선순위 큐와 heapq 를 쓰지 않고 풀어보았습니다.

 

인덱스 트리와 세그멘트 트리는 같은 의미라고 합니다. 세그멘트 트리가 학술적인 용어로도 사용되고 원래는 세그멘트 트리가 먼저 쓰여졌지만 인덱스 트리라고 불리는 이유는 알고리즘 대회를 준비하는 학생들 사이에서 노드에 인덱스를 저장해서 인덱스 트리라고 부른 것이 퍼져 인덱스 트리라고도 불린다고 합니다.

 

다음에는 펜윅트리를 이용해서 문제를 해결해보고 싶습니다.