[백준 JAVA] 2042 : 구간 합 구하기

[백준 2042] 구간 합 구하기 : https://www.acmicpc.net/problem/2042

 

---

문제 조건 정리

 

1. 어떤 N개의 수가 주어져 있다.
2. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다.
3. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다.
4. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

---

입력

 

1. 첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다.
2. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다.
3. 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다.
4. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데,
   a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고
   a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -263보다 크거나 같고, 263-1보다 작거나 같은 정수이다.

---

문제를 풀기 전

 

1. 이전에 히스토그램 문제를 풀었던 경험으로 세그먼트 트리를 구상하자.
2. a가 1인 경우, b번째 수를 c로 바꿔야할 때를 위해서,
   세그먼트 트리에서 b번째 리프 노드의 인덱스를 저장해둬야겠다고 생각함.
   이는 세그먼트 트리 초기화 과정에서 별도의 배열에 저장함.
3. 2번에서 변경한 리프노드부터 부모를 타고 올라가며 루트까지 수정할 수 있는 함수가 필요하겠다.

 

---

코드

 

package baekJoon;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class BJ2042 {

    static long[] numbers;
    static long[] segTree;
    static int[] leafNodeIdx;

    public static void main(String[] args) throws IOException {
        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer stringTokenizer = new StringTokenizer(bufferedReader.readLine(), " ");

        int n = Integer.parseInt(stringTokenizer.nextToken());
        int m = Integer.parseInt(stringTokenizer.nextToken());
        int k = Integer.parseInt(stringTokenizer.nextToken());

        numbers = new long[n + 1];
        segTree = new long[4 * n];
        leafNodeIdx = new int[n + 1];

        for (int i = 1; i <= n; i++) {
            numbers[i] = Long.parseLong(bufferedReader.readLine());
        }

        initSegTree(1, 1, n);


        for (int i = 0; i < m + k; i++) {
            stringTokenizer = new StringTokenizer(bufferedReader.readLine(), " ");
            long a = Long.parseLong(stringTokenizer.nextToken());
            int b = Integer.parseInt(stringTokenizer.nextToken());
            long c = Long.parseLong(stringTokenizer.nextToken());

            if (a == 1) {
                changeNode(b, c);
                continue;
            }

            if (a == 2) {
                System.out.println(calcSum(b, (int) c, 1, n, 1));
                continue;
            }
        }

    }

    private static void changeNode(int b, long c) {
        int targetIdx = leafNodeIdx[b]; // 바뀌어야할 노드의 인덱스 (b번째 리프 노드)
        long targetValue = segTree[targetIdx]; // 그 노드가 가진 값
        segTree[targetIdx] = c; // 리프 노드 변경

        long diff = targetValue - c; // 기존 값에서 얼마나 차이가 나는지
        for (targetIdx /= 2; targetIdx > 0; targetIdx /= 2) {
            segTree[targetIdx] -= diff; // 그 차이만큼 부모 노드 변경
        }
    }

    // start : 구간 합의 시작
    // end : 구간 합의 끝
    // left : 현재 세그먼트 트리 노드 범위의 시작
    // right : 현재 세그먼트 트리 노드 범위의 시작
    // nodeIdx : 현재 세그먼트 트리 노드 인덱스
    private static long calcSum(int start, int end, int left, int right, int nodeIdx) {
        // 현재 노드의 범위가 구간 합 범위에 완전히 속하지 않는다면
        if (left > end || right < start) {
            return 0; // 그 구간의 합은 0
        }

        // 현재 노드의 범위가 구간 합 범위에 완전히 속한다면
        if (left >= start && right <= end) {
            return segTree[nodeIdx]; // 현재 노드 값 반환(포함), 현재 노드는 현재 노드의 서브 트리의 리프 노드의 합
        }

        int mid = (left + right) / 2;
        
        // 세그먼트 트리의 현재 노드 왼쪽 자식 탐색
        long leftSum = calcSum(start, end, left, mid, nodeIdx * 2);
        // 세그먼트 트리의 현재 노드 오른쪽 자식 탐색
        long rightSum = calcSum(start, end, mid + 1, right, nodeIdx * 2 + 1);

        return leftSum + rightSum;
    }


    private static void initSegTree(int nodeIdx, int start, int end) {
        // 리프 노드
        if (start == end) {
            segTree[nodeIdx] = numbers[start];
            leafNodeIdx[start] = nodeIdx; // 세그먼트 트리의 리프 노드 인덱스를 따로 저장
            return;
        }

        int midIdx = (start + end) / 2;

        // 왼쪽 자식
        initSegTree(nodeIdx * 2, start, midIdx);

        // 오른쪽 자식
        initSegTree(nodeIdx * 2 + 1, midIdx+1, end);

        // 좌우 자식 노드 값을 더하여 부모 노드 값 갱신
        // 부모 노드 값 = 왼쪽 자식 값 + 오른쪽 자식 값
        segTree[nodeIdx] = segTree[nodeIdx * 2] + segTree[nodeIdx * 2 + 1];

    }
}

 

---

문제를 풀고 난 후

 

세그먼트 트리의 대표적인 예시인 구간 합을 위해서 연속해서 세그먼트 트리 문제를 풀어봤다.

직전에 풀었던 문제였지만, 그래도 헷갈렸고 그래서 좋았다. 연속해서 풀기를 잘했다고 생각함.. 나를 믿지 않았기에 받은 기회였다.

 

헷갈렸던 포인트는 구간 합의 범위를 쪼개고, 옮기며 탐색하는 것이 아니라, 구간 합의 범위는 고정적이며, 세그먼트 트리를 이에 맞게 구성해야하고 트리의 자식을 살피며 구간 합의 범위와 올바르게 맞춰나가야했다.

 

같은 알고리즘을 사용해도 문제를 풀어나가는 방법이 달랐다. 세그먼트 트리 알고리즘이 유용한 최적화 방법일 뿐, 문제를 해결하는 직접적인 방법이 아니라서 그런 것 같기도 하다.

 

굿