파이썬 코딩테스트 강좌, 세그먼트 트리

작성일: 2023년 10월 8일

1. 세그먼트 트리 소개

세그먼트 트리는 효율적인 구간 쿼리 처리와 업데이트를 위한 데이터 구조입니다.
주로 배열과 같은 데이터 집합에서 구간의 합, 최소값, 최대값 등을 빠르게 처리할 수 있도록 설계되었습니다.
세그먼트 트리는 완전 이진 트리 형태로 구성되어 있으며, 트리의 각 노드는 특정 구간의 정보를 저장합니다.

세그먼트 트리의 주요 특징은 다음과 같습니다:

구간 쿼리: 지정된 범위 내의 값을 빠르게 조회할 수 있습니다.
업데이트 기능: 데이터가 변경될 때마다 트리를 효율적으로 업데이트할 수 있습니다.
비교적 낮은 메모리 사용: 배열을 사용하는 것과 비교할 때 상대적으로 적은 메모리를 필요로 합니다.

2. 문제 설명

다음과 같은 문제를 고려해보겠습니다. 정수 배열 arr이 주어질 때, 특정 구간 [l, r]의 합을 구하는 쿼리와 i번째 위치의 값을 val로 업데이트하는 두 가지 기능을 지원하는 프로그램을 작성하세요.

문제의 입력 형식은 다음과 같습니다:

첫 번째 줄: 배열의 크기 N (1 ≤ N ≤ 100,000)
두 번째 줄: 배열의 원소 arr[1], arr[2], ..., arr[N]
세 번째 줄: 쿼리의 수 Q
다음 Q개의 줄: 각 쿼리를 나타내는 세 정수 type, l, r (type = 1: 구간 합 쿼리, type = 2: 업데이트 쿼리)

예를 들어, 다음과 같은 입력을 고려해볼 수 있습니다:

여기서 첫 번째 쿼리는 구간 [1, 3]의 합을 요청하고, 두 번째 쿼리는 두 번째 원소를 10으로 업데이트합니다. 세 번째 쿼리는 업데이트 후 구간 [1, 5]의 합을 요청합니다.

3. 문제 해결 과정

세그먼트 트리를 사용하여 이 문제를 해결하는 방법은 다음과 같습니다.

3.1. 세그먼트 트리 구성

우선 주어진 배열 arr로부터 세그먼트 트리를 구성해야 합니다.
부모 노드는 자식 노드들의 값을 합하여 저장합니다.
트리는 다음과 같이 초기화할 수 있습니다:

class SegmentTree:
    def __init__(self, data):
        self.n = len(data)
        self.tree = [0] * (2 * self.n)
        # 리프 노드에 데이터 넣기
        for i in range(self.n):
            self.tree[self.n + i] = data[i]
        # 내부 노드 계산
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[i * 2] + self.tree[i * 2 + 1]

3.2. 구간 합 쿼리 처리

구간 합 쿼리를 처리하기 위해서는 리프 노드에서 루트 노드까지 거슬러 올라가야 합니다.
구간 [l, r]의 합을 구하기 위해 다음과 같이 구현할 수 있습니다:

    def query(self, l, r):
        result = 0
        l += self.n
        r += self.n + 1
        while l < r:
            if l % 2 == 1:
                result += self.tree[l]
                l += 1
            if r % 2 == 1:
                r -= 1
                result += self.tree[r]
            l //= 2
            r //= 2
        return result

3.3. 업데이트 처리

업데이트 쿼리는 특정 인덱스의 값을 변경한 후, 해당 노드를 수정하고 부모 노드들에게도 영향을 미쳐야 합니다.
다음과 같이 구현할 수 있습니다:

    def update(self, index, value):
        index += self.n
        self.tree[index] = value
        while index > 1:
            index //= 2
            self.tree[index] = self.tree[index * 2] + self.tree[index * 2 + 1]

3.4. 전체 코드

이제 위의 구성 요소들을 모두 포함하는 전체 코드를 작성해봅시다:

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx]); idx += 1
    arr = [0] * N
    for i in range(N):
        arr[i] = int(data[idx]); idx += 1
    Q = int(data[idx]); idx += 1
    
    seg_tree = SegmentTree(arr)
    
    output = []
    for _ in range(Q):
        query_type = int(data[idx]); idx += 1
        l = int(data[idx]); idx += 1
        r = int(data[idx]); idx += 1
        if query_type == 1:
            result = seg_tree.query(l - 1, r - 1)
            output.append(str(result))
        elif query_type == 2:
            seg_tree.update(l - 1, r)
    
    print('\n'.join(output))

if __name__ == "__main__":
    main()

4. 시간 복잡도 분석

세그먼트 트리의 시간 복잡도는 다음과 같습니다:

세그먼트 트리의 구성: O(N)
구간 합 쿼리: O(log N)
업데이트 쿼리: O(log N)

따라서 이 알고리즘은 대규모 데이터에서도 효율적으로 동작할 수 있습니다.

5. 결론

이번 글에서는 세그먼트 트리를 사용하여 구간 합 쿼리와 업데이트 쿼리를 처리하는 방법에 대해 알아보았습니다.
세그먼트 트리는 다양한 문제에서 유용하게 사용될 수 있는 강력한 데이터 구조입니다.
코딩 테스트 시, 구간 쿼리 관련 문제가 나오면 세그먼트 트리를 고려해보는 것이 좋습니다.

목차