[백준 2042번] 구간 합 구하기

문제 개요

문제 링크 : https://www.acmicpc.net/problem/2042

 

해당 문제를 요약해보면 숫자 N개가 있고, 다음 2가지 요청을 반복해서 처리해야하는 문제입니다.

 

예를 들어 우리는 지금 숫자 카드 N장을 가지고 있다고 가정을 하겠습니다.

[1] [2] [3] [4] [5]

 

이제 다른 사람이 아래 처럼 계속 요청을 하는겁니다.

1. 3번째 카드 숫자를 6으로 바꿔주세요.
2. 2번째부터 5번째 카드 숫자를 모두 더해주세요.

만일 요청이 수십만 번 이상 반복도리 경우 그냥 for문으로 처리한다면, 시간 초과가 날 가능성이 높습니다.

그래서 빠르게 처리하는 방법이 필요합니다.

 

그걸 해결할 수 있는 방법이 세그먼트 트리입니다.

 

세그먼트 트리란?

세그먼트 트리란?

- 숫자 카드들을 트리 모양으로 묶어서 관리하는 방법

특정 구간의 합 / 최솟 값 / 최댓 값 등을 빠르게 구할 수 있도록 배열을 트리 구조로 나눈 자료 구조

용어	설명
리프 노드	배열의 한 원소를 나타내는 노드
내부 노드	자식 노드들의 구간 합을 저장
업데이트	특정 원소의 변경을 트리 전체에 반영
쿼리	구간 합을 O(log N)으로 계산

• 전체를 반으로 나눠서 왼쪽과 오른쪽으로 구간 별로 정리합니다.
• 그리고 각 구간의 합을 미리 저장해두고, 
• 누군가 구간합 알려줘 라고 하면 트리에서 필요한 정보만 찾아서 빠르게 답을 하게 됩니다.

쉽게 말해보면

상황	일반 배열	세그먼트 트리
숫자 바꾸기	O(1)	O(log N)
구간 합 구하기	O(N)	O(log N)

 

이렇게 볼 수 있습니다.

 

문제 해석

입력한 값이 아래와 같이 있습니다.

5 2 2
1
2
3
4
5
1 3 6
2 2 5
1 5 2
2 3 5

 

여기서 해석을 해보면 

• 숫자 5개: [1, 2, 3, 4, 5]
• 연산 4개 (업데이트 2번, 질의 2번)

각 요청을 살펴보면

1. 1 3 6 → 3번째 숫자를 6으로 변경
   • [1, 2, 6, 4, 5] 으로 변경
2. 2 2 5 → 2번째부터 5번째 숫자의 합을 구하라
   • 2 + 6 + 4 + 5 = 17
3. 1 5 2 → 5번째 숫자를 2로 변경
   • [1, 2, 6, 4, 2]
4. 2 3 5 → 3번째부터 5번재 합을 구하라
   • 6 + 4 + 2 = 12

이렇게 과정이 이루어지고 최종 결과는

17과 12가 나오게 됩니다.

 

이제 코드로 확인해보겠습니다.

 

import sys

# sys.stdin.readline을 사용하면 input()보다 입력 속도가 훨씬 빠릅니다.
input = sys.stdin.readline

"""
Python의 기본 재귀 깊이 제한은 1000입니다.
세그먼트 트리는 재귀적으로 노드를 생성하므로,
배열 크기가 수만 이상일 경우 'RecursionError'가 발생할 수 있습니다.
코딩 테스트 환경에서는 오류 메시지를 확인할 수 없기 때문에,
아래와 같이 재귀 깊이 제한을 넉넉하게 늘려주는 것이 안전합니다.
"""
sys.setrecursionlimit(10**6)

def solve():
    # N: 숫자 개수, M: 업데이트 횟수, K: 구간합 쿼리 횟수
    N, M, K = map(int, input().split())

    # 숫자 입력 받기 (1번 인덱스부터 사용하기 위해 앞에 0 추가)
    arr = [0] + [int(input()) for _ in range(N)]

    # 세그먼트 트리 배열 (4 * N이면 충분)
    tree = [0] * (4 * N)

    # 트리 초기화 함수
    def init(node, start, end):
        # 리프 노드: start == end
        if start == end:
            tree[node] = arr[start]
        else:
            mid = (start + end) // 2
            init(2 * node, start, mid)
            init(2 * node + 1, mid + 1, end)
            tree[node] = tree[2 * node] + tree[2 * node + 1]

    # 업데이트 함수: arr[idx] 값을 바꿨을 때, 트리 반영
    def update(node, start, end, idx, diff):
        if idx < start or idx > end:
            return  # 변경할 인덱스가 이 노드의 범위에 없으면 패스

        tree[node] += diff  # 이 노드가 커버하는 구간의 합에 diff 반영

        if start != end:  # 내부 노드라면 자식도 갱신해야 함
            mid = (start + end) // 2
            update(2 * node, start, mid, idx, diff)
            update(2 * node + 1, mid + 1, end, idx, diff)

    # 구간 합 쿼리 함수: [left, right] 범위의 합을 구함
    def query(node, start, end, left, right):
        # 이 노드의 구간이 요청 범위와 겹치지 않음
        if end < left or start > right:
            return 0

        # 이 노드의 구간이 요청 범위에 완전히 포함됨
        if left <= start and end <= right:
            return tree[node]

        # 일부만 겹친 경우 → 자식 노드에서 합을 구해서 더함
        mid = (start + end) // 2
        left_sum = query(2 * node, start, mid, left, right)
        right_sum = query(2 * node + 1, mid + 1, end, left, right)
        return left_sum + right_sum

    # 1단계: 세그먼트 트리 초기화
    init(1, 1, N)

    # 2단계: M + K 개의 연산 처리
    for _ in range(M + K):
        cmd = list(map(int, input().split()))
        if cmd[0] == 1:
            # 업데이트: idx 위치의 값을 new_val로 변경
            _, idx, new_val = cmd
            diff = new_val - arr[idx]  # 기존 값과의 차이
            arr[idx] = new_val
            update(1, 1, N, idx, diff)
        else:
            # 쿼리: [left, right] 구간 합 출력
            _, left, right = cmd
            print(query(1, 1, N, left, right))

if __name__ == "__main__":
    solve()

 

위 코드의 실행 흐름을 다시 보겠습니다.

초기 상태 : [1, 2, 3, 4, 5]

명령	동작	결과
1 3 6	3번 값을 3 → 6으로 바꿈	트리 갱신됨
2 2 5	2~5 합: 2+6+4+5 = 17	출력
1 5 2	5번 값을 5 → 2로 바꿈	트리 갱신됨
2 3 5	3~5 합: 6+4+2 = 12	출력