백준 2042: 구간 합 구하기 [골드 Ⅰ/ Java]

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

2042번: 구간 합 구하기

---

 

 

 

Solution

세그먼트 트리, 펜윅 트리

 

 

배열의 부분 합을 구할 때 보통 합 배열을 이용한다.

하지만 크기가 7인 배열의 3번째 값이 바뀌면, 합 배열을 3번부터 끝까지 다 수정해야 한다.

 

 

 

대신에 세그먼트 트리로 배열을 이진 트리처럼 저장하면, 관련된 부분만 수정할 수 있다.

 

합 배열 (좌) / 세그먼트 트리 (우)

 

 

자료 구조만 바꾸었는데 탐색 횟수가 반으로 줄었음을 알 수 있다.

이 효과는 배열 크기가 커질수록 유용할 것이다. 

 

 

세그먼트 트리로 값이 자주 바뀌는 배열의 부분 합/곱, 최솟값/최댓값을 효율적으로 구할 수 있다.

 

 

 

 

 

세그먼트 트리를 생성해 보자.

하지만 사실은 새로운 1차원 배열을 생성하는 것이다.

 

본 배열 값들은 leaf node 에 저장한다.

그리고 그 부모 노드들을 채워가면서 이진 트리를 완성한다.

 

 

 

이때 리프 노드 개수는, 원래 배열의 크기(7) 보다 크면서 가장 작은 2의 제곱수(8) 이어야 한다.

그럼 전체 트리의 노드 개수는 15 (1 + 2¹ + 2² + 2³) 개 이다.

 

 

 

원래 배열을 새로운 배열에 적용하려면 인덱싱을 해주어야 한다.

 

원래 배열에 2³ 을 더해준 인덱스에 값을 넣어준다. 

그럼 1번 인덱스까지 넣어주고, 리프 노드는 2 ³ 부터 시작한다.

 

그럼 8번 노드, 9번 노드의 합은 4번 노드에 넣어주고, 4번 노드와 5번 노드의 합은 2번 노드에 넣어준다.

 

자식 노드 2n 와 2n+1 의 부모 노드는 n 이다.

 

 

 

 

 

세그먼트 트리로 인덱스 9 부터 13 까지 구간 합을 구해보자.

 

 

시작 노드는 9, 끝 노드는 13 이고,

시작 노드가 짝수 / 홀수 일 때와 끝 노드가 짝수 / 홀수 일 때 경우가 나뉜다.

 

 

시작 노드가

짝수이면 부모 노드에 포함되므로 따로 합을 구하지 않아도 되고

홀수이면 부모 노드에 포함되지 않으므로 값을 따로 더해주어야 한다.

 

끝 노드가

짝수이면 부모 노드에 포함되지 않으므로  값을 따로 더해주어야 하고,

홀수이면 부모 노드에 포함되므로 따로 합을 구하지 않아도 된다.

 

 

 

그리고 시작 노드는 + 1, 끝 노드는 - 1을 해준다.

 

 

 

 

 

Code

 

 

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

public class Main {

    static long[] tree;
    static StringBuilder sb = new StringBuilder();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int M = Integer.parseInt(st.nextToken()); // 변경 횟수
        int K = Integer.parseInt(st.nextToken()); // 연산 횟수

        int lev = 0;
        while ((1 << lev) < N) lev++; // 9: 2⁴
        int treeSize = (int) Math.pow(2, lev + 1); // 0, 인덱스 1부터 시작 ~ 16-31 = 32개
        int compl = treeSize / 2; // 1+15, 16

        tree = new long[treeSize]; // 0 ~ 2^5-1
        for (int i = compl; i < compl + N; i++)
            tree[i] = Long.parseLong(br.readLine());

        setTree(compl + N - 1); // treeSize-1

        for (int i = 0; i < M + K; i++) {
            st = new StringTokenizer(br.readLine());
            long a = Integer.parseInt(st.nextToken());
            int s = Integer.parseInt(st.nextToken());
            long e = Long.parseLong(st.nextToken());

            // 변경
            if (a == 1) changeVal(compl - 1 + s, e);
            else if (a == 2) { // 구간 합
                s = s - 1 + compl;
                e = e - 1 + compl;

                sb.append(getSum(s, (int) e)).append("\n");
            } else break;
        }

        System.out.print(sb);
    }

    private static void changeVal(int index, long val) {
        long diff = val - tree[index];
        while (index > 0) {  // 루트노드까지 부모노드
            tree[index] += diff;
            index /= 2;
        }
    }

    private static long getSum(int s, int e) {
        long partSum = 0;

        while (s <= e) {
            if (s % 2 == 1) {  // odd 계산
                partSum += tree[s];
                s++;
            }
            if (e % 2 == 0) {  // even 계산
                partSum += tree[e];
                e--;
            }
            s /= 2;
            e /= 2;
        }

        return partSum;
    }

    private static void setTree(int i) {
        while (i != 1) {
            tree[i / 2] += tree[i];  // 8, 9 -> 4
            i--;
        }
    }

}